Upload folder using huggingface_hub

main.py

@@ -1,6 +1,8 @@
-from fastapi import FastAPI, Request, Form
+from fastapi import FastAPI, Request, Form, HTTPException
 from fastapi.templating import Jinja2Templates
 from fastapi.responses import HTMLResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 import torch
 from transformers import AutoTokenizer
 from optimum.onnxruntime import ORTModelForSequenceClassification
@@ -14,36 +16,44 @@ import time
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
+app = FastAPI(title="Spam Detection System API")
+
+# 1. Enable CORS for cross-platform system calls
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], # In production, replace with your actual platform URL
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 templates = Jinja2Templates(directory="templates")
 
-# Model Paths (Local Project Directory)
+# Model Paths
 QUANT_MINILM_PATH = "models/minilm_int8"
 QUANT_XLM_PATH = "models/xlm_roberta_int8"
 ML_MODEL_PATH = "models/baseline/spam_detection_model.pkl"
 
+# Request Schema for System Calls
+class PredictionRequest(BaseModel):
+    text: str
+    model: str = "minilm" # Default to the fastest production model
+
 # Load Models and Tokenizers
 logger.info("Loading Quantized Production Models...")
-
-# 1. MiniLM INT8
 minilm_tokenizer = AutoTokenizer.from_pretrained(QUANT_MINILM_PATH)
 minilm_model = ORTModelForSequenceClassification.from_pretrained(QUANT_MINILM_PATH, file_name="model_quantized.onnx")
 
-# 2. XLM-Roberta INT8
 xlm_tokenizer = AutoTokenizer.from_pretrained(QUANT_XLM_PATH)
 xlm_model = ORTModelForSequenceClassification.from_pretrained(QUANT_XLM_PATH, file_name="model_quantized.onnx")
 
-# 3. ML Model (Scikit-learn)
 ml_model = joblib.load(ML_MODEL_PATH)
 
 def preprocess_ml(text):
-    if not isinstance(text, str):
-        text = str(text)
-    text = text.lower()
+    text = str(text).lower()
     text = re.sub(r'[^\w\s]', '', text)
     text = re.sub(r'\d+', '', text)
-    text = text.strip()
-    return text
+    return text.strip()
 
 def predict_ort(text, model, tokenizer):
     start_time = time.time()
@@ -52,12 +62,43 @@ def predict_ort(text, model, tokenizer):
         outputs = model(**inputs)
     probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
     prediction = torch.argmax(probs, dim=-1).item()
-    
-    latency = (time.time() - start_time) * 1000 # ms
-    
+    latency = (time.time() - start_time) * 1000
     label = model.config.id2label[prediction] if hasattr(model.config, "id2label") else ("Spam" if prediction == 1 else "Ham")
-    label = label.title()
-    return label, probs[0][prediction].item(), latency
+    return label.title(), probs[0][prediction].item(), latency
+
+# --- SYSTEM ENDPOINTS ---
+
+@app.get("/health")
+async def health_check():
+    """System health check to verify models are loaded."""
+    return {"status": "ready", "models_loaded": ["minilm_int8", "xlm_roberta_int8", "baseline_ml"]}
+
+@app.post("/api/v1/predict")
+async def system_predict(request: PredictionRequest):
+    """Formal JSON API endpoint for system-to-system integration."""
+    if not request.text:
+        raise HTTPException(status_code=400, detail="Text is required")
+    
+    if request.model == "minilm":
+        label, conf, lat = predict_ort(request.text, minilm_model, minilm_tokenizer)
+    elif request.model == "xlm":
+        label, conf, lat = predict_ort(request.text, xlm_model, xlm_tokenizer)
+    else:
+        # Fallback to ML Model
+        start = time.time()
+        pred = ml_model.predict([preprocess_ml(request.text)])[0]
+        lat = (time.time() - start) * 1000
+        label = "Spam" if pred == 1 else "Ham"
+        conf = 1.0
+        
+    return {
+        "label": label,
+        "confidence": float(conf),
+        "latency_ms": float(lat),
+        "model_used": request.model
+    }
+
+# --- WEB UI ENDPOINTS ---
 
 @app.get("/", response_class=HTMLResponse)
 async def read_item(request: Request):
@@ -65,28 +106,17 @@ async def read_item(request: Request):
 
 @app.post("/compare")
 async def compare_models(text: str = Form(...)):
+    # Legacy endpoint for the HTML frontend
     results = []
-    
-    # MiniLM INT8
     l1, s1, t1 = predict_ort(text, minilm_model, minilm_tokenizer)
     results.append({"model": "MiniLM v2 (INT8)", "label": l1, "confidence": f"{s1:.4f}", "latency": f"{t1:.2f}ms"})
-    
-    # XLM-Roberta INT8
     l2, s2, t2 = predict_ort(text, xlm_model, xlm_tokenizer)
     results.append({"model": "XLM-R v2 (INT8)", "label": l2, "confidence": f"{s2:.4f}", "latency": f"{t2:.2f}ms"})
-    
-    # ML Model
     start3 = time.time()
-    clean_text = preprocess_ml(text)
-    pred3 = ml_model.predict([clean_text])[0]
+    pred3 = ml_model.predict([preprocess_ml(text)])[0]
     t3 = (time.time() - start3) * 1000
     l3 = "Spam" if pred3 == 1 else "Ham"
-    try:
-        p3 = ml_model.predict_proba([clean_text])[0][pred3]
-    except:
-        p3 = 1.0
-    results.append({"model": "ML Model (English)", "label": l3, "confidence": f"{p3:.4f}", "latency": f"{t3:.2f}ms"})
-    
+    results.append({"model": "ML Model (English)", "label": l3, "confidence": "1.0000", "latency": f"{t3:.2f}ms"})
     return results
 
 if __name__ == "__main__":