| | """ |
| | Inference script for UnixCoder-MIL |
| | ===================================== |
| | Usage: Simply run this script with your code samples |
| | """ |
| |
|
| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | from transformers import AutoTokenizer, AutoModel, AutoConfig, AutoModelForSequenceClassification |
| | from safetensors.torch import load_file |
| | import numpy as np |
| |
|
| | DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
| | CLASS_NAMES = ["Human", "AI-Generated", "Hybrid", "Adversarial"] |
| |
|
| | class MilUnixCoder(nn.Module): |
| | def __init__(self, model_name="microsoft/unixcoder-base", chunk_size=512, stride=256, max_chunks=16): |
| | super().__init__() |
| | self.config = AutoConfig.from_pretrained(model_name) |
| | self.unixcoder = AutoModel.from_pretrained(model_name) |
| | self.chunk_size, self.stride, self.max_chunks = chunk_size, stride, max_chunks |
| | self.classifier = nn.Linear(self.config.hidden_size, 4) |
| | self.dropout = nn.Dropout(0.1) |
| | def forward(self, input_ids, attention_mask=None): |
| | B, L = input_ids.size() |
| | if attention_mask is None: attention_mask = torch.ones_like(input_ids) |
| | if L > self.chunk_size: |
| | c_ids = input_ids.unfold(1, self.chunk_size, self.stride) |
| | c_mask = attention_mask.unfold(1, self.chunk_size, self.stride) |
| | nc = min(c_ids.size(1), self.max_chunks) |
| | flat_ids = c_ids[:,:nc,:].contiguous().view(-1, self.chunk_size) |
| | flat_mask = c_mask[:,:nc,:].contiguous().view(-1, self.chunk_size) |
| | else: |
| | nc, flat_ids, flat_mask = 1, input_ids, attention_mask |
| | out = self.unixcoder(input_ids=flat_ids, attention_mask=flat_mask) |
| | logits = self.classifier(self.dropout(out.last_hidden_state[:, 0, :])) |
| | return torch.max(logits.view(B, nc, -1), dim=1)[0] |
| |
|
| | def load_model(): |
| | """Load the model and tokenizer""" |
| | from huggingface_hub import hf_hub_download |
| | repo = "YoungDSMLKZ/UnixCoder-MIL" |
| | tokenizer = AutoTokenizer.from_pretrained(repo) |
| | model = MilUnixCoder("microsoft/unixcoder-base") |
| | weights_path = hf_hub_download(repo_id=repo, filename="model.safetensors") |
| | model.load_state_dict(load_file(weights_path)) |
| | model.to(DEVICE).eval() |
| | return model, tokenizer |
| |
|
| | def predict(code: str, model, tokenizer) -> dict: |
| | """Predict class for a single code sample""" |
| | inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=4096, padding=True).to(DEVICE) |
| | with torch.no_grad(): |
| | logits = model(inputs["input_ids"], inputs["attention_mask"]) |
| | probs = F.softmax(logits, dim=-1)[0] |
| | pred = torch.argmax(probs).item() |
| | return {"class": CLASS_NAMES[pred], "confidence": probs[pred].item()} |
| |
|
| | if __name__ == "__main__": |
| | print("Loading model...") |
| | model, tokenizer = load_model() |
| | |
| | |
| | test_code = """ |
| | def hello_world(): |
| | print("Hello, World!") |
| | """ |
| | |
| | result = predict(test_code, model, tokenizer) |
| | print(f"Predicted: {result['class']} (confidence: {result['confidence']:.2%})") |
| |
|
