Update ai_text_detector_valid_final.py

ai_text_detector_valid_final.py

@@ -1,90 +1,131 @@
+import os
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import requests
 import numpy as np
+import re
 
-# Multiple AI text detection models
-MODELS = {    
-    "DeBERTa Detector": "distilbert-base-uncased-finetuned-sst-2-english",
-    "MonkeyDAnh":"MonkeyDAnh/my-awesome-ai-detector-roberta-base-v4-human-vs-machine-finetune",
-    "Andreas122001":"andreas122001/roberta-academic-detector"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ---------------------------
+# ModernBERT Models (SzegedAI Workflow)
+# ---------------------------
+model1_path = "modernbert.bin"
+model2_path = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed12"
+model3_path = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed22"
+
+tokenizer_modernbert = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
+
+model_1 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
+model_1.load_state_dict(torch.load(model1_path, map_location=device))
+model_1.to(device).eval()
+
+model_2 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
+model_2.load_state_dict(torch.hub.load_state_dict_from_url(model2_path, map_location=device))
+model_2.to(device).eval()
+
+model_3 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
+model_3.load_state_dict(torch.hub.load_state_dict_from_url(model3_path, map_location=device))
+model_3.to(device).eval()
+
+label_mapping = {
+    0: '13B', 1: '30B', 2: '65B', 3: '7B', 4: 'GLM130B', 5: 'bloom_7b',
+    6: 'bloomz', 7: 'cohere', 8: 'davinci', 9: 'dolly', 10: 'dolly-v2-12b',
+    11: 'flan_t5_base', 12: 'flan_t5_large', 13: 'flan_t5_small',
+    14: 'flan_t5_xl', 15: 'flan_t5_xxl', 16: 'gemma-7b-it', 17: 'gemma2-9b-it',
+    18: 'gpt-3.5-turbo', 19: 'gpt-35', 20: 'gpt4', 21: 'gpt4o',
+    22: 'gpt_j', 23: 'gpt_neox', 24: 'human', 25: 'llama3-70b', 26: 'llama3-8b',
+    27: 'mixtral-8x7b', 28: 'opt_1.3b', 29: 'opt_125m', 30: 'opt_13b',
+    31: 'opt_2.7b', 32: 'opt_30b', 33: 'opt_350m', 34: 'opt_6.7b',
+    35: 'opt_iml_30b', 36: 'opt_iml_max_1.3b', 37: 't0_11b', 38: 't0_3b',
+    39: 'text-davinci-002', 40: 'text-davinci-003'
 }
-    
-def load_model(model_id):
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForSequenceClassification.from_pretrained(model_id)
-    return tokenizer, model
-
-def predict(text, tokenizer, model):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+
+def clean_text(text: str) -> str:
+    text = re.sub(r"\s{2,}", " ", text)
+    text = re.sub(r"\s+([,.;:?!])", r"\1", text)
+    return text
+
+def classify_szegedai(text: str):
+    """ModernBERT ensemble detector (replaces SzegedAI Space call)."""
+    cleaned_text = clean_text(text)
+    if not cleaned_text.strip():
+        return {"error": "Empty text"}
+
+    inputs = tokenizer_modernbert(cleaned_text, return_tensors="pt", truncation=True, padding=True).to(device)
+
     with torch.no_grad():
-        outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    return probs[0].numpy()  # [human_prob, ai_prob]
-
-
-def verdict(ai_prob):
-    """Return a human-readable verdict based on AI probability"""
-    if ai_prob < 20:
-        return "Most likely human-written."
-    elif 20 <= ai_prob < 40:
-        return "Possibly human-written with minimal AI assistance."
-    elif 40 <= ai_prob < 60:
-        return "Unclear – could be either human or AI-assisted."
-    elif 60 <= ai_prob < 80:
-        return "Possibly AI-generated, or a human using AI assistance."
-    else:  # ai_prob >= 80
-        return "Likely AI-generated or heavily AI-assisted."
+        logits_1 = model_1(**inputs).logits
+        logits_2 = model_2(**inputs).logits
+        logits_3 = model_3(**inputs).logits
 
-def detect_text(text):
-    results = {}
-    ai_scores = []
+        probs = (torch.softmax(logits_1, dim=1) +
+                 torch.softmax(logits_2, dim=1) +
+                 torch.softmax(logits_3, dim=1)) / 3
 
-    for name, model_id in MODELS.items():
-        try:
-            tokenizer, model = load_model(model_id)
-            probs = predict(text, tokenizer, model)
-            human_score, ai_score = probs
-            results[name] = {
-                "Human Probability": round(float(human_score) * 100, 2),
-                "AI Probability": round(float(ai_score) * 100, 2),
-            }
-            ai_scores.append(ai_score)
-        except Exception as e:
-            results[name] = {"error": str(e)}
-
-    # ------------------ Final Score (Average) ------------------
+    probs = probs[0]
+    ai_probs = probs.clone()
+    ai_probs[24] = 0  # "human" label index
+    ai_total_prob = ai_probs.sum().item() * 100
+    human_prob = 100 - ai_total_prob
+
+    ai_index = torch.argmax(ai_probs).item()
+    ai_model = label_mapping[ai_index]
+
+    return {
+        "Human Probability": round(human_prob, 2),
+        "AI Probability": round(ai_total_prob, 2),
+        "Identified LLM": ai_model
+    }
+
+# ---------------------------
+# Your Other Detectors
+# ---------------------------
+MODELS = {
+    "DeBERTa Detector": "distilbert-base-uncased-finetuned-sst-2-english",
+    "MonkeyDAnh": "MonkeyDAnh/my-awesome-ai-detector-roberta-base-v4-human-vs-machine-finetune",
+    "Andreas122001": "andreas122001/roberta-academic-detector",
+}
+
+def run_hf_model(model_id, text):
     try:
-        ai_scores, human_scores = [], []
-
-        for r in results.values():
-            if isinstance(r, dict) and "AI Probability" in r and "Human Probability" in r:
-                ai_scores.append(r["AI Probability"])
-                human_scores.append(r["Human Probability"])
-
-        if ai_scores and human_scores:
-            avg_ai = sum(ai_scores) / len(ai_scores)
-            avg_human = sum(human_scores) / len(human_scores)
-
-            results["Final Score"] = {
-                # "Human Probability (average)": float(round(avg_human, 2)),
-                # "AI Probability (average)": float(round(avg_ai, 2))
-                # "Verdict": verdict(avg_ai)
-                verdict(avg_ai)
-            }
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForSequenceClassification.from_pretrained(model_id)
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+        with torch.no_grad():
+            logits = model(**inputs).logits
+        probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
+        return {"Human Probability": float(probs[0]*100), "AI Probability": float(probs[1]*100)}
     except Exception as e:
-        results["Final Score"] = {"error": str(e)}
+        return {"error": str(e)}
 
-    return results
+# ---------------------------
+# Main Detector
+# ---------------------------
+def detect_text(text):
+    results = {}
+    # HuggingFace transformer models
+    for name, model_id in MODELS.items():
+        results[name] = run_hf_model(model_id, text)
+
+    # SzegedAI ModernBERT ensemble
+    results["SzegedAI Detector"] = classify_szegedai(text)
 
+    # Final verdict
+    ai_probs = []
+    for v in results.values():
+        if "AI Probability" in v:
+            ai_probs.append(v["AI Probability"])
+    avg_ai = np.mean(ai_probs) if ai_probs else 0
+    if avg_ai > 80:
+        verdict = "Likely AI-generated"
+    elif avg_ai > 40:
+        verdict = "Possibly human-written with AI assistance"
+    else:
+        verdict = "Likely human-written"
+    results["Final Score"] = {"Verdict": verdict}
+    return results
 
 if __name__ == "__main__":
-    text = input("Enter text to analyze:\n")
-    output = detect_text(text)
-    print("\n--- Detection Results ---")
-    for model, scores in output.items():
-        print(f"\n[{model}]")
-        for k, v in scores.items():
-            if isinstance(v, (int, float)):  # only add % for numeric values
-                print(f"{k}: {v}%")
-            else:
-                print(f"{k}: {v}")
+    sample = "This is a test sentence written by AI or human."
+    print(detect_text(sample))