File size: 1,583 Bytes
9f48851 f7a530f 9f48851 f7a530f 9f48851 f7a530f 9f48851 d1fe2cb f7a530f d1fe2cb f7a530f d1fe2cb 8e14057 |
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 51 52 |
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import RobertaTokenizerFast, RobertaForSequenceClassification, TextClassificationPipeline
import uvicorn
# Define FastAPI app
app = FastAPI()
# Load Model on Startup
HUGGINGFACE_MODEL_PATH = "bespin-global/klue-roberta-small-3i4k-intent-classification"
print("Loading model...") # Log message
try:
loaded_tokenizer = RobertaTokenizerFast.from_pretrained(HUGGINGFACE_MODEL_PATH)
loaded_model = RobertaForSequenceClassification.from_pretrained(HUGGINGFACE_MODEL_PATH)
# Create Text Classification Pipeline
text_classifier = TextClassificationPipeline(
tokenizer=loaded_tokenizer,
model=loaded_model,
return_all_scores=True
)
print("Model loaded successfully.") # Log message
except Exception as e:
print(f"Error loading model: {e}")
text_classifier = None
# Health Check Endpoint
@app.get("/")
def hello():
return {"Message": "Space is running Good.", "Status": "Healthy"}
# Define Pydantic Model for Input Validation
class PredictionRequest(BaseModel):
sentence: str
# Prediction Endpoint
@app.post("/predict")
def predict_intent(request: PredictionRequest):
if text_classifier is None:
return {"error": "Model not found"}
sentence = request.sentence.strip() # Correct way to get JSON input
preds_list = text_classifier(sentence)
best_pred = max(preds_list[0], key=lambda x: x["score"]) # Get highest-scoring intent
return {"predicted_intent": best_pred["label"], "confidence": best_pred["score"]}
|