File size: 797 Bytes
5a56923 ae8bbe8 0c886b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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 model once when the app starts
classifier = pipeline("zero-shot-classification", model="syedkhalid076/DeBERTa-Zero-Shot-Classification")
# Request body model
class ClassificationRequest(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"]}
|