solved batch problem

app.py

@@ -2,25 +2,17 @@ import pickle
 import os
 from fastapi import FastAPI
 from pydantic import BaseModel
+from typing import List
 
-# -----------------------------
-# FastAPI initialization
-# -----------------------------
 app = FastAPI(
     title="RTL Log Severity Classifier",
-    description="Machine learning API that predicts severity of RTL verification logs.",
-    version="1.0"
+    description="Batch severity prediction for RTL verification logs",
+    version="1.1"
 )
 
-# -----------------------------
-# Model paths
-# -----------------------------
 VECTORIZER_PATH = "vectorizer.pkl"
 MODEL_PATH = "severity_model.pkl"
 
-# -----------------------------
-# Severity mapping
-# -----------------------------
 REVERSE_MAP = {
     0: "INFO",
     1: "WARNING",
@@ -28,57 +20,79 @@ REVERSE_MAP = {
     3: "CRITICAL"
 }
 
-# -----------------------------
-# Load artifacts safely
-# -----------------------------
-if not os.path.exists(VECTORIZER_PATH):
-    raise RuntimeError("vectorizer.pkl not found")
-
-if not os.path.exists(MODEL_PATH):
-    raise RuntimeError("severity_model.pkl not found")
-
+# Load artifacts
 with open(VECTORIZER_PATH, "rb") as f:
     vectorizer = pickle.load(f)
 
 with open(MODEL_PATH, "rb") as f:
     model = pickle.load(f)
 
-# -----------------------------
-# Request schema
-# -----------------------------
-class LogRequest(BaseModel):
+
+# ---------- Request Schemas ----------
+
+class LogItem(BaseModel):
     module: str
     message: str
 
 
-# -----------------------------
-# Health check
-# -----------------------------
+class BatchRequest(BaseModel):
+    logs: List[LogItem]
+
+
+# ---------- Health ----------
+
 @app.get("/")
-def health_check():
+def health():
     return {
         "status": "running",
         "model": "RTL Severity Classifier",
-        "classes": ["INFO", "WARNING", "ERROR", "CRITICAL"]
+        "batch_support": True
     }
 
 
-# -----------------------------
-# Prediction endpoint
-# -----------------------------
+# ---------- Single Prediction ----------
+
 @app.post("/predict")
-def predict_severity(request: LogRequest):
+def predict(log: LogItem):
+
+    text = log.module + " " + log.message
+
+    vec = vectorizer.transform([text])
+
+    pred = model.predict(vec)[0]
+
+    return {
+        "module": log.module,
+        "message": log.message,
+        "predicted_severity": REVERSE_MAP[pred]
+    }
+
+
+# ---------- Batch Prediction ----------
+
+@app.post("/predict_batch")
+def predict_batch(request: BatchRequest):
+
+    texts = [
+        log.module + " " + log.message
+        for log in request.logs
+    ]
+
+    vectors = vectorizer.transform(texts)
 
-    text = request.module + " " + request.message
+    preds = model.predict(vectors)
 
-    vector = vectorizer.transform([text])
+    results = []
 
-    pred = model.predict(vector)[0]
+    for i, p in enumerate(preds):
 
-    severity = REVERSE_MAP[pred]
+        results.append({
+            "module": request.logs[i].module,
+            "message": request.logs[i].message,
+            "predicted_severity": REVERSE_MAP[int(p)]
+        })
 
     return {
-        "module": request.module,
-        "message": request.message,
-        "predicted_severity": severity
+        "count": len(results),
+        "results": results
     }