Update app.py

app.py

@@ -1,116 +1,220 @@
 import os
+import re
+import shutil
 import torch
-from fastapi import FastAPI, WebSocket, UploadFile, File
-from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI
+from pydantic import BaseModel
+
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from huggingface_hub import hf_hub_download
-
-# ==========================================================
-# ✅ Environment Setup (Safe Caching)
-# ==========================================================
-os.environ["HF_HOME"] = "/tmp/huggingface"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
-os.environ["HF_HUB_CACHE"] = "/tmp/huggingface"
-os.makedirs("/tmp/huggingface", exist_ok=True)
-os.makedirs("/tmp/models", exist_ok=True)
-
-# ==========================================================
-# ✅ Model Paths
-# ==========================================================
+from tokenizers.normalizers import Sequence, Replace, Strip
+from tokenizers import Regex
+import atexit
+
+# =====================================================
+# ✅ Safe Cache Setup (Runtime)
+# =====================================================
+CACHE_DIR = "/tmp/huggingface"
+
+# Clear any old cache to prevent exceeding 50G
+if os.path.exists(CACHE_DIR):
+    shutil.rmtree(CACHE_DIR, ignore_errors=True)
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+# Set environment paths
+os.environ.update({
+    "HF_HOME": CACHE_DIR,
+    "TRANSFORMERS_CACHE": CACHE_DIR,
+    "HF_DATASETS_CACHE": CACHE_DIR,
+    "HF_HUB_CACHE": CACHE_DIR,
+    "TORCH_HOME": CACHE_DIR,
+    "XDG_CACHE_HOME": CACHE_DIR,
+    "TORCHINDUCTOR_CACHE_DIR": CACHE_DIR,
+    "TORCH_LOGS": "off"
+})
+
+# Auto cleanup on shutdown
+@atexit.register
+def cleanup_cache():
+    shutil.rmtree(CACHE_DIR, ignore_errors=True)
+
+# =====================================================
+# ✅ Model Setup
+# =====================================================
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
+
+# --- Model paths ---
 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"
-
-# ==========================================================
-# ✅ Helper: Load Model (Local or Remote)
-# ==========================================================
-def load_model(path_or_url, model_name="modernbert"):
-    try:
-        if path_or_url.startswith("http"):
-            print(f"Downloading model from {path_or_url}...")
-            filename = os.path.join("/tmp/models", os.path.basename(path_or_url))
-            if not os.path.exists(filename):
-                torch.hub.download_url_to_file(path_or_url, filename)
-            path_or_url = filename
-
-        tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-        model = AutoModelForSequenceClassification.from_pretrained(
-            "bert-base-uncased",
-            state_dict=torch.load(path_or_url, map_location="cpu"),
+model2_url = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed12"
+model3_url = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed22"
+
+def load_model(base_path=None, url=None):
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "answerdotai/ModernBERT-base", num_labels=41
+    )
+    if url:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url, map_location=device, progress=False, check_hash=False
+
+
+
+
+
+
+
+
+
         )
-        model.eval()
-        print(f"✅ Loaded model: {path_or_url}")
-        return tokenizer, model
-    except Exception as e:
-        print(f"❌ Failed to load model from {path_or_url}: {e}")
-        return None, None
-
-# ==========================================================
-# ✅ Load All Models
-# ==========================================================
-tokenizer1, model1 = load_model(model1_path, "model1")
-tokenizer2, model2 = load_model(model2_path, "model2")
-tokenizer3, model3 = load_model(model3_path, "model3")
-
-# ==========================================================
+    else:
+        state_dict = torch.load(base_path, map_location=device)
+    model.load_state_dict(state_dict)
+    model.to(device).eval()
+    return model
+
+model_1 = load_model(model1_path)
+model_2 = load_model(url=model2_url)
+model_3 = load_model(url=model3_url)
+
+# =====================================================
+# ✅ Label Mapping
+# =====================================================
+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'
+}
+
+# =====================================================
+# ✅ Text Cleaning & Tokenizer Normalization
+# =====================================================
+def clean_text(text: str) -> str:
+    text = re.sub(r'\s{2,}', ' ', text)
+    text = re.sub(r'\s+([,.;:?!])', r'\1', text)
+    return text
+
+newline_to_space = Replace(Regex(r'\s*\n\s*'), " ")
+join_hyphen_break = Replace(Regex(r'(\w+)[--]\s*\n\s*(\w+)'), r"\1\2")
+tokenizer.backend_tokenizer.normalizer = Sequence([
+    tokenizer.backend_tokenizer.normalizer,
+    join_hyphen_break,
+    newline_to_space,
+    Strip()
+])
+
+# =====================================================
+# ✅ Analysis Logic
+# =====================================================
+def analyze_text_block(text: str):
+    cleaned_text = clean_text(text)
+    inputs = tokenizer(cleaned_text, return_tensors="pt", truncation=True, padding=True).to(device)
+
+    with torch.no_grad():
+        logits_1 = model_1(**inputs).logits
+        logits_2 = model_2(**inputs).logits
+        logits_3 = model_3(**inputs).logits
+
+        avg_probs = (
+            torch.softmax(logits_1, dim=1) +
+            torch.softmax(logits_2, dim=1) +
+            torch.softmax(logits_3, dim=1)
+        ) / 3
+
+        probs = avg_probs[0]
+        human_prob = probs[24].item()
+        ai_probs = probs.clone()
+        ai_probs[24] = 0
+        ai_total_prob = ai_probs.sum().item()
+
+        total = human_prob + ai_total_prob
+        human_percentage = (human_prob / total) * 100
+        ai_percentage = (ai_total_prob / total) * 100
+        ai_model_index = torch.argmax(ai_probs).item()
+
+        return {
+            "human_written_score": round(human_percentage / 100, 4),
+            "ai_generated_score": round(ai_percentage / 100, 4),
+            "predicted_model": label_mapping[ai_model_index]
+        }
+
+def split_into_paragraphs(text: str):
+    return [p.strip() for p in re.split(r'\n\s*\n', text.strip()) if p.strip()]
+
+# =====================================================
 # ✅ FastAPI Setup
-# ==========================================================
-app = FastAPI(title="ModernBERT AI Detection API")
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"], allow_credentials=True,
-    allow_methods=["*"], allow_headers=["*"]
-)
-
-# ==========================================================
-# ✅ Root Route
-# ==========================================================
-@app.get("/")
-def home():
-    return {"message": "ModernBERT API is running"}
-
-# ==========================================================
-# ✅ Text Inference Endpoint
-# ==========================================================
+# =====================================================
+app = FastAPI(title="ModernBERT AI Text Detector")
+
+class InputText(BaseModel):
+    text: str
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+
+
+
+
+
+
+
+
+
 @app.post("/analyze")
-async def analyze_text(data: dict):
-    text = data.get("text", "")
-    if not text.strip():
-        return JSONResponse({"error": "No text provided"}, status_code=400)
+async def analyze(data: InputText):
+    text = data.text.strip()
+    if not text:
+        return {"success": False, "code": 400, "message": "Empty input text"}
 
-    inputs = tokenizer1(text, return_tensors="pt", truncation=True, padding=True)
-    results = {}
+    total_words = len(text.split())
+    full_result = analyze_text_block(text)
+    fake_percentage = round(full_result["ai_generated_score"] * 100, 2)
+    ai_words = int(total_words * (fake_percentage / 100))
+    results = []
 
-    with torch.no_grad():
-        if model1:
-            results["model1"] = torch.softmax(model1(**inputs).logits, dim=-1).tolist()[0]
-        if model2:
-            results["model2"] = torch.softmax(model2(**inputs).logits, dim=-1).tolist()[0]
-        if model3:
-            results["model3"] = torch.softmax(model3(**inputs).logits, dim=-1).tolist()[0]
-
-    avg_score = torch.tensor([results[m][0] for m in results]).mean().item()
-    return {"results": results, "avg_score": avg_score}
-
-# ==========================================================
-# ✅ WebSocket for Real-Time Scoring
-# ==========================================================
-@app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-    try:
-        while True:
-            data = await websocket.receive_text()
-            inputs = tokenizer1(data, return_tensors="pt", truncation=True, padding=True)
-            with torch.no_grad():
-                output = torch.softmax(model1(**inputs).logits, dim=-1)
-            await websocket.send_json({"score": output.tolist()[0]})
-    except Exception:
-        await websocket.close()
-
-# ==========================================================
-# ✅ Run Command
-# ==========================================================
-# uvicorn app:app --host 0.0.0.0 --
+    if fake_percentage > 50:
+        paragraphs = split_into_paragraphs(text)
+        ai_words, total_words = 0, 0
+        for p in paragraphs:
+            res = analyze_text_block(p)
+            wc = len(p.split())
+            total_words += wc
+            ai_words += wc * res["ai_generated_score"]
+            results.append({
+                "paragraph": p,
+                "ai_generated_score": res["ai_generated_score"],
+                "human_written_score": res["human_written_score"],
+                "predicted_model": res["predicted_model"]
+            })
+        fake_percentage = round((ai_words / total_words) * 100, 2)
+
+    feedback = (
+        "Most of Your Text is AI/GPT Generated"
+        if fake_percentage > 50
+        else "Most of Your Text Appears Human-Written"
+    )
+
+    return {
+        "success": True,
+        "code": 200,
+        "message": "analysis completed",
+        "data": {
+            "fakePercentage": fake_percentage,
+            "isHuman": round(100 - fake_percentage, 2),
+            "textWords": total_words,
+            "aiWords": ai_words,
+            "paragraphs": results,
+            "predicted_model": full_result["predicted_model"],
+            "feedback": feedback,
+            "input_text": text,
+            "detected_language": "en"
+        }
+    }