service-internal's picture
Update main.py
7aaed34 verified
raw
history blame contribute delete
814 Bytes
import os
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
# Define FastAPI app
app = FastAPI()
# Load zero-shot classification pipeline
classifier = pipeline("zero-shot-classification", model="syedkhalid076/DeBERTa-Zero-Shot-Classification")
# Define input schema
class InputText(BaseModel):
text: str
@app.post("/classify")
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]
}