File size: 821 Bytes
8cf8359 0c886b1 8cf8359 2083d34 0c886b1 2083d34 0c886b1 2083d34 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
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="MoritzLaurer/deberta-v3-large-zeroshot-v2.0"
)
# 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]
} |