SourceCodeAuthorCheck-SLM-10M

A ~10 million parameter Small Language Model (SLM) Transformer designed for binary classification to detect whether Python source code was written by a human or generated by AI.

The model was trained on a dataset of human-written code extracted from GitHub (Q3 2017 and prior) and synthetic AI-generated code mimicking Q3 2026 generative AI paradigms.

Live Demo

You can test the model interactively without writing any code by visiting our Gradio Web UI Space.

Usage

To use this model locally, you must include the model architecture class in your script before loading the weights.

1. Define the Architecture

import torch
import torch.nn as nn
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download

class SourceCodeAuthorCheck(nn.Module):
    def __init__(self, vocab_size=50257, d_model=128, nhead=8, num_layers=4, dim_feedforward=512):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = nn.Parameter(torch.zeros(1, 1024, d_model))

        encoder_layers = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layers, num_layers=num_layers)
        self.fc = nn.Linear(d_model, 1)

    def forward(self, input_ids, attention_mask):
        seq_len = input_ids.size(1)
        x = self.embedding(input_ids) + self.pos_encoder[:, :seq_len, :]

        src_key_padding_mask = ~attention_mask.bool()
        x = self.transformer(x, src_key_padding_mask=src_key_padding_mask)

        mask_expanded = attention_mask.unsqueeze(-1).float()
        sum_embeddings = torch.sum(x * mask_expanded, 1)
        sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)
        pooled = sum_embeddings / sum_mask

        return self.fc(pooled)

2. Load Weights and Run Inference

Once the class is defined, you can automatically pull the weights from Hugging Face and run inference cleanly using this helper function:

# Initialize device and tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

# Download and load model weights from the Hugging Face Hub
model = SourceCodeAuthorCheck().to(device)
model_path = hf_hub_download(repo_id="assix-research/SourceCodeAuthorCheck-SLM-10M", filename="source_code_classifier.pth")
model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
model.eval()

def predict_source(code_snippet):
    """Predicts whether a given Python snippet is human-written or AI-generated."""
    inputs = tokenizer(
        code_snippet, 
        return_tensors="pt", 
        truncation=True, 
        padding="max_length", 
        max_length=1024
    ).to(device)

    with torch.no_grad():
        if torch.cuda.is_available():
            with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
                logits = model(inputs['input_ids'], inputs['attention_mask'])
        else:
            logits = model(inputs['input_ids'], inputs['attention_mask'])

        prob = torch.sigmoid(logits).item()

    verdict = "AI Generated" if prob > 0.5 else "Human Written"
    print(f"Verdict: {verdict} (AI Probability: {prob:.1%})")

# --- Example Usage ---
sample_code = """
def calculate_factorial(n):
    if n == 0:
        return 1
    return n * calculate_factorial(n-1)
"""

predict_source(sample_code)