Update main.py

main.py

@@ -1,23 +1,26 @@
-import os
-# ✅ Set cache to writable directory
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
-os.environ["HF_HOME"] = "/tmp/hf-home"
-
 from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from transformers import pipeline
 
-app = FastAPI()
+# Load zero-shot classification pipeline
+classifier = pipeline(
+    "zero-shot-classification",
+    model="MoritzLaurer/deberta-v3-large-zeroshot-v2.0"
+)
 
-# Load model once when the app starts
-classifier = pipeline("zero-shot-classification", model="syedkhalid076/DeBERTa-Zero-Shot-Classification")
+# Define FastAPI app
+app = FastAPI()
 
-# Request body model
-class ClassificationRequest(BaseModel):
+# Define input schema
+class InputText(BaseModel):
     text: str
 
 @app.post("/classify")
-async def classify(request: ClassificationRequest):
-    candidate_labels = ["contains electronic components", "doesn't contain electronic components"]
-    result = classifier(request.text, candidate_labels, multi_label=False)
-    return {"labels": result["labels"], "scores": result["scores"]}
+async def classify_text(data: InputText):
+    candidate_labels = ["contains electronic components", "does not contain electronic components"]
+    result = classifier(data.text, candidate_labels, multi_label=False)
+    return {
+        "input": data.text,
+        "label": result["labels"][0],
+        "score": result["scores"][0]
+    }