Update app.py

app.py

@@ -1,80 +1,49 @@
 import torch
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-import logging
-import json
-import os
-
-# Set up logging configuration
-logging.basicConfig(level=logging.INFO)
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 # Initialize the FastAPI app
 app = FastAPI()
 
-# Load the trained model (adjust the path to your saved model)
-model = torch.load("path/to/your/model.pth", map_location=torch.device("cpu"))  # Replace with your actual model path
+# 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):
-    """
-    Data model for the prediction request.
-    
-    Attributes:
-        text (str): Input text for model inference.
-    """
     text: str
 
 class PredictionResponse(BaseModel):
-    """
-    Data model for the prediction response.
-    
-    Attributes:
-        text (str): The original input text.
-        prediction (str): The predicted result from the model.
-    """
     text: str
     prediction: str
 
 # Define prediction endpoint
 @app.post("/predict", response_model=PredictionResponse)
 async def predict(request: PredictionRequest):
-    """
-    Endpoint for generating a prediction based on input text.
-    
-    Args:
-        request (PredictionRequest): The request body containing the input text.
-        
-    Returns:
-        PredictionResponse: The response body containing the original text and prediction.
-    
-    Raises:
-        HTTPException: If any error occurs during the prediction process.
-    """
     try:
-        # Tokenize the input text (assuming you're using a tokenizer for text inputs)
-        inputs = tokenizer(request.text, return_tensors="pt")
+        # Tokenize the input text
+        inputs = tokenizer(request.text, return_tensors="pt", truncation=True, padding=True)
         
         # Perform inference with the model
-        outputs = model(**inputs)
+        with torch.no_grad():
+            outputs = model(**inputs)
+        
+        # Get the predicted class
+        prediction = torch.argmax(outputs.logits, dim=-1).item()
         
-        # Get the predicted token and decode it back to text
-        prediction = tokenizer.decode(outputs.logits.argmax(-1)[0], skip_special_tokens=True)
+        # 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=prediction)
+        return PredictionResponse(text=request.text, prediction=predicted_label)
     except Exception as e:
-        logging.error("Error during prediction", exc_info=True)
         raise HTTPException(status_code=500, detail="Prediction failed")
 
 # Define health check endpoint
 @app.get("/health")
 async def health_check():
-    """
-    Health check endpoint to verify if the service is up and running.
-    
-    Returns:
-        dict: A dictionary containing the status of the service.
-    """
-    logging.info("Health check requested.")
     return {"status": "healthy"}