Spaces:
Runtime error
Runtime error
| import os | |
| import sys | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" | |
| from pathlib import Path | |
| from fastapi import FastAPI, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| from contextlib import asynccontextmanager | |
| from typing import Optional | |
| import torch | |
| import pickle | |
| import re | |
| import redis | |
| from huggingface_hub import hf_hub_download | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| # ── Constants ──────────────────────────────────────────── | |
| CONFIDENCE_THRESHOLD = 10.0 | |
| OOS_INTENT = "oos" | |
| OOS_RESPONSE = "I'm not sure I understand that. Could you rephrase or ask me something banking related?" | |
| MAX_LEN = 64 | |
| DEVICE = torch.device("cpu") | |
| MODEL_REPO = "RexTheSlasher/intent-classifier-bert" | |
| # ── Redis Cloud ────────────────────────────────────────── | |
| r = redis.Redis( | |
| host=os.getenv('REDIS_HOST'), | |
| port=int(os.getenv('REDIS_PORT', 6379)), | |
| password=os.getenv('REDIS_PASSWORD'), | |
| decode_responses=True | |
| ) | |
| # ── Load model from Hugging Face Hub ──────────────────── | |
| print("Loading model from Hugging Face...") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO) | |
| model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO) | |
| model = model.to(DEVICE) | |
| model.eval() | |
| # ── Load label encoder from Hugging Face Hub ──────────── | |
| encoder_path = hf_hub_download(repo_id=MODEL_REPO, filename="label_encoder.pkl") | |
| with open(encoder_path, "rb") as f: | |
| label_encoder = pickle.load(f) | |
| print(f"Model loaded! Classes: {len(label_encoder.classes_)}") | |
| # ── Text cleaning ──────────────────────────────────────── | |
| def clean_text(text: str) -> str: | |
| text = text.lower() | |
| text = re.sub(r"[^a-z0-9\s]", "", text) | |
| text = re.sub(r"\s+", " ", text).strip() | |
| return text | |
| # ── Predict ────────────────────────────────────────────── | |
| def predict(text: str) -> dict: | |
| cleaned = clean_text(text) | |
| encoding = tokenizer( | |
| cleaned, | |
| max_length=MAX_LEN, | |
| padding="max_length", | |
| truncation=True, | |
| return_tensors="pt" | |
| ) | |
| input_ids = encoding["input_ids"].to(DEVICE) | |
| attention_mask = encoding["attention_mask"].to(DEVICE) | |
| with torch.no_grad(): | |
| outputs = model(input_ids=input_ids, attention_mask=attention_mask) | |
| probabilities = torch.softmax(outputs.logits, dim=1) | |
| confidence, predicted_class = torch.max(probabilities, dim=1) | |
| predicted_label = label_encoder.inverse_transform([predicted_class.item()])[0] | |
| confidence_score = round(confidence.item() * 100, 2) | |
| top3_probs, top3_indices = torch.topk(probabilities, 3, dim=1) | |
| top3 = [ | |
| { | |
| "intent": label_encoder.inverse_transform([idx.item()])[0], | |
| "confidence": round(prob.item() * 100, 2) | |
| } | |
| for prob, idx in zip(top3_probs[0], top3_indices[0]) | |
| ] | |
| is_low_confidence = confidence_score < CONFIDENCE_THRESHOLD | |
| if predicted_label == OOS_INTENT or is_low_confidence: | |
| predicted_label = "out_of_scope" | |
| is_oos = True | |
| else: | |
| is_oos = False | |
| return { | |
| "input_text": text, | |
| "cleaned_text": cleaned, | |
| "intent": predicted_label, | |
| "confidence": confidence_score, | |
| "top3": top3, | |
| "is_low_confidence": is_low_confidence, | |
| "is_oos": is_oos, | |
| "response": OOS_RESPONSE if is_oos else f"I detected your intent as {predicted_label}" | |
| } | |
| # ── FastAPI App ────────────────────────────────────────── | |
| async def lifespan(app: FastAPI): | |
| print("✅ API started!") | |
| yield | |
| app = FastAPI( | |
| title="Intent Classification Chatbot API", | |
| version="2.0.0", | |
| lifespan=lifespan | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"] | |
| ) | |
| # ── Models ─────────────────────────────────────────────── | |
| class PredictRequest(BaseModel): | |
| text: str | |
| session_id: Optional[str] = None | |
| class IntentResult(BaseModel): | |
| intent: str | |
| confidence: float | |
| class PredictResponse(BaseModel): | |
| session_id: str | |
| input_text: str | |
| cleaned_text: str | |
| intent: str | |
| confidence: float | |
| top3: list[IntentResult] | |
| last_intent: Optional[str] = None | |
| is_low_confidence: bool = False | |
| is_oos: bool = False | |
| response: str = "" | |
| # ── Routes ─────────────────────────────────────────────── | |
| def root(): | |
| return {"message": "Intent Chatbot API", "status": "running", "intents": len(label_encoder.classes_)} | |
| def health(): | |
| return {"status": "healthy", "model": "bert-base-uncased", "intents": len(label_encoder.classes_)} | |
| def predict_intent(request: PredictRequest): | |
| text = request.text.strip() | |
| if not text: | |
| raise HTTPException(status_code=400, detail="Text cannot be empty") | |
| if len(text) < 2: | |
| raise HTTPException(status_code=400, detail="Text too short") | |
| if len(text) > 200: | |
| raise HTTPException(status_code=400, detail="Text too long") | |
| import uuid, json | |
| from datetime import datetime | |
| session_id = request.session_id or str(uuid.uuid4()) | |
| # Get last intent | |
| existing = r.get(f"session:{session_id}") | |
| messages = json.loads(existing) if existing else [] | |
| last_intent = None | |
| for msg in reversed(messages): | |
| if msg.get("intent"): | |
| last_intent = msg["intent"] | |
| break | |
| result = predict(text) | |
| # Save to Redis | |
| messages.append({"role": "user", "text": text, "timestamp": datetime.now().isoformat()}) | |
| messages.append({"role": "bot", "text": result["intent"], "intent": result["intent"], "confidence": result["confidence"], "timestamp": datetime.now().isoformat()}) | |
| r.setex(f"session:{session_id}", 3600, json.dumps(messages)) | |
| return PredictResponse( | |
| session_id=session_id, | |
| input_text=result["input_text"], | |
| cleaned_text=result["cleaned_text"], | |
| intent=result["intent"], | |
| confidence=result["confidence"], | |
| top3=[IntentResult(**r_) for r_ in result["top3"]], | |
| last_intent=last_intent, | |
| is_low_confidence=result["is_low_confidence"], | |
| is_oos=result["is_oos"], | |
| response=result["response"] | |
| ) | |
| def get_history(session_id: str): | |
| existing = r.get(f"session:{session_id}") | |
| if not existing: | |
| raise HTTPException(status_code=404, detail="Session not found") | |
| return {"session_id": session_id, "history": json.loads(existing)} | |
| def delete_session(session_id: str): | |
| r.delete(f"session:{session_id}") | |
| return {"message": "Session cleared"} | |
| def get_intents(): | |
| return {"total": len(label_encoder.classes_), "intents": list(label_encoder.classes_)} | |