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]
    }