update from kimi2

app.py

@@ -1,85 +1,95 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import pipeline
-import mysql.connector
+"""
+FastAPI server inside Hugging Face Space
+POST /predict  ->  zero-shot subject prediction + save to TiDB
+"""
 import os
+import time
+from contextlib import asynccontextmanager
 
-os.environ['HF_HOME'] = '/tmp/huggingface'
+import mysql.connector
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 
-app = FastAPI()
+# ---------- load model ONCE ----------
+MODEL_NAME = "MoritzLaurer/deberta-v3-large-zeroshot-v1.1-all-33"
+LABELS = [
+    "Mathematics", "Physics", "Chemistry", "Biology",
+    "History", "Geography", "Literature", "Computer-Science"
+]
 
-# Database connection settings
-DB_HOST = "gateway01.ap-southeast-1.prod.aws.tidbcloud.com"
-DB_PORT = 4000
-DB_USER = "4V44XYoMA7okY9v.root"
-DB_PASS = "aW2CrSwcTgjFhNAb"
-DB_NAME = "final_project"
+ml_models = {}
 
-# Create MySQL connection
-conn = mysql.connector.connect(
-    host=DB_HOST,
-    port=DB_PORT,
-    user=DB_USER,
-    password=DB_PASS,
-    database=DB_NAME,
-    ssl_verify_cert=True,
-    ssl_verify_identity=True
-)
-cursor = conn.cursor()
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # load at startup
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+    model.eval()
+    if torch.cuda.is_available():
+        model.cuda()
+    ml_models["tokenizer"] = tokenizer
+    ml_models["model"] = model
+    yield
+    # shutdown
+    ml_models.clear()
 
-# Load model
-classifier = pipeline(
-    "zero-shot-classification",
-    model="MoritzLaurer/deberta-v3-large-zeroshot-v1.1-all-33"
-)
+app = FastAPI(lifespan=lifespan)
 
-# Labels
-subject_labels = [
-    "Physics", "Chemistry", "Biology", "Astronomy",
-    "Earth Science", "Environmental Science",
-    "Algebra", "Geometry", "Calculus", "Statistics",
-    "Probability", "Number Theory",
-    "English Language", "English Literature",
-    "Tamil Language", "Tamil Literature",
-    "History", "Geography", "Political Science", "Economics",
-    "Sociology", "Psychology", "Philosophy",
-    "Computer Science", "Data Science", "Artificial Intelligence",
-    "Robotics", "Biotechnology", "Engineering",
-    "Fine Arts", "Music", "Dance", "Theater",
-    "Business Studies", "Accountancy", "Entrepreneurship",
-    "Physical Education", "Health Science"
-]
+# ---------- DB helper ----------
+def get_conn():
+    return mysql.connector.connect(
+        host=os.getenv("DB_HOST"),
+        port=int(os.getenv("DB_PORT", 4000)),
+        user=os.getenv("DB_USER"),
+        password=os.getenv("DB_PASS"),
+        database=os.getenv("DB_NAME"),
+        ssl_ca=os.getenv("DB_SSL_CA_PATH") or None
+    )
 
-# Request model
-class TextInput(BaseModel):
+# ---------- request schema ----------
+class PredictRequest(BaseModel):
     student_id: str
     text: str
 
+# ---------- API endpoint ----------
 @app.post("/predict")
-def predict_topic(data: TextInput):
-    # Predict subject
-    result = classifier(
-        data.text,
-        candidate_labels=subject_labels,
-        hypothesis_template="This text is about {}."
+def predict(req: PredictRequest):
+    if not req.text.strip():
+        raise HTTPException(400, "Empty text")
+    tok = ml_models["tokenizer"](
+        req.text,
+        padding=True,
+        truncation=True,
+        return_tensors="pt"
     )
-    predicted_subject = result["labels"][0]
+    if torch.cuda.is_available():
+        tok = {k: v.cuda() for k, v in tok.items()}
+    with torch.no_grad():
+        logits = ml_models["model"](**tok).logits
+        probs = torch.softmax(logits, dim=-1)[0]
+        idx = int(torch.argmax(probs))
+        subject = LABELS[idx]
 
-    # Get first 100 characters of the text
-    sample_text = data.text[:100]
+    # save to DB
+    try:
+        conn = get_conn()
+        cur = conn.cursor()
+        cur.execute(
+            "INSERT INTO predictions (student_id, text, subject, created_at) "
+            "VALUES (%s, %s, %s, %s)",
+            (req.student_id, req.text, subject, time.strftime('%Y-%m-%d %H:%M:%S'))
+        )
+        conn.commit()
+        cur.close()
+        conn.close()
+    except Exception as e:
+        print("DB error:", e)
+        raise HTTPException(500, "DB write failed")
 
-    # Save to DB
-    cursor.execute(
-        """
-        INSERT INTO log_table (student_id, input_sample, subject)
-        VALUES (%s, %s, %s)
-        """,
-        (data.student_id, sample_text, predicted_subject)
-    )
-    conn.commit()
+    return {"subject": subject}
 
-    return {
-        "student_id": data.student_id,
-        "predicted_subject": predicted_subject,
-        "sample_text": sample_text
-    }
+@app.get("/")
+def root():
+    return {"message": "Subject predictor is running"}