Spaces:
Running
Running
| from fastapi import FastAPI, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| from typing import List | |
| from transformers import pipeline | |
| import json | |
| import uvicorn | |
| import os | |
| app = FastAPI(title="The Algorithm - Cloud GPU API") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=False, # 🛡️ Sentinel: Credentials must be False when using wildcard origins to prevent CSRF | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Initialize model once on boot | |
| print("Loading XLM-RoBERTa Sentiment Pipeline into GPU...") | |
| model_name = "cardiffnlp/twitter-xlm-roberta-base-sentiment" | |
| # 🛡️ Sentinel: Pin model revision for supply chain integrity | |
| REVISION = "f2f1202" | |
| # device=0 targets the first available GPU (Lightning Studio T4/L4) | |
| sentiment_pipeline = pipeline("sentiment-analysis", model=model_name, device=0, revision=REVISION) | |
| print("Model loaded successfully.") | |
| class TextPayload(BaseModel): | |
| texts: List[str] | |
| async def analyze_sentiment(payload: TextPayload): | |
| """ | |
| Accepts a JSON list of strings. | |
| Processes them through the GPU-accelerated NLP model. | |
| Returns a JSON list of integer sentiment scores. | |
| """ | |
| try: | |
| if not payload.texts: | |
| return {"scores": []} | |
| # 🛡️ Sentinel: Limit payload size to prevent RAM/GPU exhaustion (DoS) | |
| if len(payload.texts) > 2000: | |
| raise HTTPException(status_code=413, detail="Payload too large. Maximum 2000 texts allowed.") | |
| # Robust generator to stream data to the GPU in batches | |
| # This fixes the "must be type str" error and maximizes efficiency | |
| def data_gen(): | |
| for text in payload.texts: | |
| yield str(text) if text else "" | |
| # Inference via pipeline using the generator | |
| results = sentiment_pipeline(data_gen(), batch_size=128, truncation=True) | |
| sentiment_scores = [] | |
| for r in results: | |
| label = r['label'].lower() | |
| if 'positive' in label or label == 'label_2': | |
| sentiment_scores.append(1) | |
| elif 'negative' in label or label == 'label_0': | |
| sentiment_scores.append(-1) | |
| else: | |
| sentiment_scores.append(0) | |
| return {"scores": sentiment_scores} | |
| except Exception as e: | |
| # 🛡️ Sentinel: Mask internal exceptions to prevent information disclosure | |
| print(f"Error during GPU inference: {str(e)}") | |
| raise HTTPException(status_code=500, detail="An internal error occurred during GPU inference.") | |
| async def root(): | |
| return {"message": "The Algorithm Cloud API is running successfully. Please post to /analyze for sentiment scoring."} | |
| async def health_check(): | |
| return {"status": "online", "gpu_enabled": True} | |
| if __name__ == "__main__": | |
| # 🛡️ Sentinel: Use environment variable for host binding to follow security best practices | |
| host = os.environ.get("HOST", "127.0.0.1") | |
| uvicorn.run(app, host=host, port=8000) | |