Create app.py

app.py

@@ -0,0 +1,92 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+# 1. Model Architecture
+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. Device and Loading Initialization
+# Note: HF Free Spaces use CPU. We check for CUDA in case you upgrade the Space hardware.
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token
+
+model = SourceCodeAuthorCheck().to(device)
+
+# Download weights securely from your model repository
+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()
+
+# 3. Inference Logic
+def predict_author(code_snippet):
+    if not code_snippet or not code_snippet.strip():
+        return "Please paste valid code.", "0.0%"
+        
+    inputs = tokenizer(
+        code_snippet, 
+        return_tensors="pt", 
+        truncation=True, 
+        padding="max_length", 
+        max_length=1024
+    ).to(device)
+    
+    with torch.no_grad():
+        # Handle mixed precision safely based on available hardware
+        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()
+
+    score = round(prob * 100, 2)
+    verdict = "🤖 AI Generated" if prob > 0.5 else "👨‍💻 Human Written"
+    
+    return verdict, f"{score}%"
+
+# 4. Gradio Interface Construction
+demo = gr.Interface(
+    fn=predict_author,
+    inputs=gr.Code(language="python", label="Paste Python Source Code"),
+    outputs=[
+        gr.Textbox(label="Verdict"),
+        gr.Textbox(label="AI Probability Score")
+    ],
+    title="SourceCodeAuthorCheck SLM (10M)",
+    description="Analyze Python snippets to determine if they were written by a human or generated by an AI model.",
+    examples=[
+        ["def calculate_tax(gross_salary, deduction):\n    return gross_salary - deduction"],
+        ["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}"]
+    ]
+)
+
+if __name__ == "__main__":
+    demo.launch()