zero shot

app.py

@@ -1,24 +1,38 @@
-from fastapi import FastAPI
-import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from typing import List
+from transformers import pipeline
 
-checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
-
-sequences = ["I've been waiting for a HuggingFace course my whole life.", "So have I!"]
-
-tokens = tokenizer(sequences, padding=True, truncation=True, return_tensors="pt")
-
-# Perform inference without gradient tracking
-with torch.no_grad():
-    output = model(**tokens)
-
-# Convert logits to a list for JSON serialization
-logits = output.logits.tolist()
+# Initialize the zero-shot classification pipeline
+classifier = pipeline("zero-shot-classification")
 
+# Define the FastAPI application
 app = FastAPI()
 
+# Pydantic model for input validation
+class ClassificationRequest(BaseModel):
+    text: str = Field(..., example="This is a course about the Transformers library")
+    labels: List[str] = Field(..., example=["education", "politics", "technology"])
+
 @app.get("/")
 def greet_json():
-    return {"Hello": logits}
+    """
+    A simple GET endpoint that returns a greeting message.
+    """
+    return {"Hello": "World!"}
+
+@app.post("/classify")
+def zero_shot_classification(request: ClassificationRequest):
+    """
+    A POST endpoint that performs zero-shot classification on the input text
+    using the provided candidate labels.
+    """
+    try:
+        # Perform zero-shot classification
+        result = classifier(
+            request.text,
+            candidate_labels=request.labels
+        )
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

requirements.txt

@@ -1,4 +1,5 @@
 fastapi
 uvicorn[standard]
 torch
-transformers
+transformers
+pydantic