Spaces:
Running
Running
import os | |
from fastapi import FastAPI | |
from transformers import pipeline | |
# Set a custom Hugging Face cache directory | |
os.environ["TRANSFORMERS_CACHE"] = "/app/.cache" | |
app = FastAPI() | |
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli") | |
async def predict(data: dict): | |
labels = ["Mathematics", "Language Arts", "Social Studies", "Science"] | |
result = classifier(data["text"], labels) | |
return {"label": result["labels"][0]} | |