File size: 1,709 Bytes
ed06dec 3e66b84 177d7f0 3e66b84 ed06dec 3e66b84 177d7f0 ed06dec 3e66b84 ed06dec 3e66b84 ed06dec 3e66b84 177d7f0 ed06dec 177d7f0 ed06dec 177d7f0 ed06dec 177d7f0 3e66b84 ed06dec 3e66b84 |
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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# Initialize the FastAPI app
app = FastAPI()
# Load the model and tokenizer from Hugging Face
model_name = "Canstralian/RabbitRedux" # Replace with your model's name
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()
# Define the input and output format for prediction requests
class PredictionRequest(BaseModel):
text: str
class PredictionResponse(BaseModel):
text: str
prediction: str
# Define prediction endpoint
@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
try:
# Tokenize the input text
inputs = tokenizer(request.text, return_tensors="pt", truncation=True, padding=True)
# Perform inference with the model
with torch.no_grad():
outputs = model(**inputs)
# Get the predicted class
prediction = torch.argmax(outputs.logits, dim=-1).item()
# Map the prediction to a label (adjust as per your model's labels)
labels = ["Label 1", "Label 2", "Label 3"] # Replace with your actual labels
predicted_label = labels[prediction]
# Return the prediction response
return PredictionResponse(text=request.text, prediction=predicted_label)
except Exception as e:
raise HTTPException(status_code=500, detail="Prediction failed")
# Define health check endpoint
@app.get("/health")
async def health_check():
return {"status": "healthy"}
|