CodeDetect-UniXCoder-T2

This model is a fine-tuned version of microsoft/unixcoder-base

2. Load the Model

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_NAME = "azherali/CodeDetect-UniXCoder-T2"
MAX_LENGTH = 512

# Select device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device:", device)

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# Configure class labels
model.config.id2label = {
                "0": "human",
                "1": "qwen",
                "2": "phi",
                "3": "meta-llama",
                "4": "01-ai",
                "5": "gemma",
                "6": "ibm-granite",
                "7": "bigcode",
                "8": "deepseek",
                "9": "gpt",
                "10": "mistral"
            }
model.config.label2id = {
                "human": 0,
                "qwen": 1,
                "phi": 2,
                "meta-llama": 3,
                "01-ai": 4,
                "gemma": 5,
                "ibm-granite": 6,
                "bigcode": 7,
                "deepseek": 8,
                "gpt": 9,
                "mistral": 10
                }

model.to(device)
model.eval()

print("Number of classes:", model.config.num_labels)
print("Label mapping:", model.config.id2label)

3. Prediction Function

The predict_code() function takes source code as input and returns the predicted label, class ID, and probability scores.

def predict_code(code: str) -> dict:
    if not isinstance(code, str) or not code.strip():
        raise ValueError("Please provide a non-empty code string.")

    # Tokenize source code
    inputs = tokenizer(
        code,
        truncation=True,
        max_length=MAX_LENGTH,
        padding=True,
        return_tensors="pt",
    )

    # Move inputs to the selected device
    inputs = {
        key: value.to(device)
        for key, value in inputs.items()
    }

    # Perform inference
    with torch.inference_mode():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=-1)[0]

    # Get predicted class
    predicted_id = int(probabilities.argmax().item())

    labels = model.config.id2label or {}

    predicted_label = labels.get(
        predicted_id,
        labels.get(str(predicted_id), f"class_{predicted_id}")
    )

    # Return prediction and class probabilities
    return {
        "predicted_class_id": predicted_id,
        "predicted_label": predicted_label,
        "class_probabilities": {
            str(labels.get(i, labels.get(str(i), f"class_{i}"))):
                round(float(probabilities[i].item()), 4)
            for i in range(len(probabilities))
        },
    }

4. Example Inference

code = """
def add(a, b):
    return a + b
"""

result = predict_code(code)
print(result)

5. Example Output

The following output is illustrative. Actual predictions and scores depend on the input and model behavior.

{
        "predicted_class_id": 9,
        "predicted_label": "gpt",
        "class_probabilities": {
            "human": 0.0001,
            "qwen": 0.4001,
            "phi": 0.0014,
            "meta-llama": 0.0405,
            "01-ai": 0.0024,
            "gemma": 0.0007,
            "ibm-granite": 0.0034,
            "bigcode": 0.0,
            "deepseek": 0.0162,
            "gpt": 0.5345,
            "mistral": 0.0007,
        },
    }

6. Output Description

Training hyperparameters

The following hyperparameters were used during training:

• learning_rate: 2e-05
• train_batch_size: 64
• eval_batch_size: 64
• seed: 42
• optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
• lr_scheduler_type: linear
• lr_scheduler_warmup_steps: 0.1
• num_epochs: 5
• mixed_precision_training: Native AMP

Training results

Framework versions

• Transformers 5.13.1
• Pytorch 2.11.0+cu130
• Datasets 4.0.0
• Tokenizers 0.22.2