Spaces:
Sleeping
Sleeping
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 | |
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] | |
} |