assix-research/SourceCodeAuthorCheck-UI

1import gradio as gr2import torch3import torch.nn as nn4from transformers import AutoTokenizer5from huggingface_hub import hf_hub_download6import spaces7 8# 1. Model Architecture9class SourceCodeAuthorCheck(nn.Module):10    def __init__(self, vocab_size=50257, d_model=128, nhead=8, num_layers=4, dim_feedforward=512):11        super().__init__()12        self.embedding = nn.Embedding(vocab_size, d_model)13        self.pos_encoder = nn.Parameter(torch.zeros(1, 1024, d_model))14        15        encoder_layers = nn.TransformerEncoderLayer(16            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True17        )18        self.transformer = nn.TransformerEncoder(encoder_layers, num_layers=num_layers)19        self.fc = nn.Linear(d_model, 1)20 21    def forward(self, input_ids, attention_mask):22        seq_len = input_ids.size(1)23        x = self.embedding(input_ids) + self.pos_encoder[:, :seq_len, :]24        25        src_key_padding_mask = ~attention_mask.bool()26        x = self.transformer(x, src_key_padding_mask=src_key_padding_mask)27        28        mask_expanded = attention_mask.unsqueeze(-1).float()29        sum_embeddings = torch.sum(x * mask_expanded, 1)30        sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)31        pooled = sum_embeddings / sum_mask32        33        return self.fc(pooled)34 35# 2. Device and Loading Initialization36tokenizer = AutoTokenizer.from_pretrained("gpt2")37tokenizer.pad_token = tokenizer.eos_token38 39# Load model globally on CPU first40model = SourceCodeAuthorCheck()41model_path = hf_hub_download(repo_id="assix-research/SourceCodeAuthorCheck-SLM-10M", filename="source_code_classifier.pth")42model.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=True))43model.eval()44 45# 3. Inference Logic with ZeroGPU Decorator46@spaces.GPU47def predict_author(code_snippet):48    if not code_snippet or not code_snippet.strip():49        return "Please paste valid code.", "0.0%"50    51    # ZeroGPU dynamically provides CUDA access inside this decorated function52    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")53    model.to(device)54        55    inputs = tokenizer(56        code_snippet, 57        return_tensors="pt", 58        truncation=True, 59        padding="max_length", 60        max_length=102461    ).to(device)62    63    with torch.no_grad():64        if torch.cuda.is_available():65            with torch.autocast(device_type='cuda', dtype=torch.bfloat16):66                logits = model(inputs['input_ids'], inputs['attention_mask'])67        else:68            logits = model(inputs['input_ids'], inputs['attention_mask'])69            70        prob = torch.sigmoid(logits).item()71 72    score = round(prob * 100, 2)73    verdict = "🤖 AI Generated" if prob > 0.5 else "👨‍💻 Human Written"74    75    # Move model back to CPU to free up ZeroGPU vRAM for other users76    model.to("cpu")77    78    return verdict, f"{score}%"79 80# 4. Gradio Interface Construction81demo = gr.Interface(82    fn=predict_author,83    inputs=gr.Code(language="python", label="Paste Python Source Code"),84    outputs=[85        gr.Textbox(label="Verdict"),86        gr.Textbox(label="AI Probability Score")87    ],88    title="SourceCodeAuthorCheck SLM (10M)",89    description="Analyze Python snippets to determine if they were written by a human or generated by an AI model.",90    examples=[91        ["def calculate_tax(gross_salary, deduction):\n    return gross_salary - deduction"],92        ["def process_data_stream_0(data_input: list[dict], strict_validation: bool = True) -> dict:\n    if not data_input:\n        return {'status': 'error', 'message': 'Empty stream'}\n    processed_results = []\n    for idx, item in enumerate(data_input):\n        transformed = {k: str(v).strip().lower() for k, v in item.items()}\n        transformed['_internal_id'] = f'gen_id_0_{idx}'\n        processed_results.append(transformed)\n    return {'status': 'success', 'data': processed_results}"]93    ]94)95 96if __name__ == "__main__":97    demo.launch()