CodeDetect-UniXCoder-T1

2. Load the Model

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_NAME = "azherali/CodeDetect-UniXCoder-T1"
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: "AI"
}

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': 0,
    'predicted_label': 'Human',
    'class_probabilities': {
        'Human': 0.9231,
        'AI': 0.0769
    }
}

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