service-internal's picture
Update main.py
2083d34 verified
raw
history blame
719 Bytes
from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import pipeline
# Load zero-shot classification pipeline
classifier = pipeline(
"zero-shot-classification",
model="MoritzLaurer/deberta-v3-large-zeroshot-v2.0"
)
# Define FastAPI app
app = FastAPI()
# 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]
}