Add Query Classification

app.py

@@ -1,4 +1,3 @@
-
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -7,6 +6,7 @@ from transformers import pipeline
 from duckduckgo_search import DDGS
 from typing import Optional, List, Dict, Any
 import logging
+import re
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -14,12 +14,12 @@ logger = logging.getLogger(__name__)
 
 # Initialize FastAPI app
 app = FastAPI(
-    title="Open Source Chat API",
-    description="A fully open source alternative to HF API using local models",
-    version="1.0.0"
+    title="Enhanced Chat API with Dynamic Model Selection",
+    description="API with query classification and dynamic model loading for QA/Summarization",
+    version="2.0.0"
 )
 
-# Configure CORS for Android app access
+# Configure CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -36,175 +36,200 @@ class ChatRequest(BaseModel):
 
 class ChatResponse(BaseModel):
     response: str
-    search_results: Optional[List[Dict[str, Any]]] = None  # Contains web search results if use_search is enabled
+    task_type: str
+    search_results: Optional[List[Dict[str, Any]]] = None
 
 class SearchRequest(BaseModel):
     query: str
     max_results: int = 5
 
+# Global pipelines with lazy loading
+classifier_pipeline = None
 qa_pipeline = None
+summarization_pipeline = None
 ner_pipeline = None
 
-def load_ner_model():
-    """Load Named Entity Recognition model"""
-    global ner_pipeline
+def load_classifier():
+    """Load zero-shot classification model"""
+    global classifier_pipeline
     try:
-        # Check if GPU is available
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device: {device}")
-
-        ner_pipeline = pipeline(
-            "ner",
-            model="dbmdz/bert-large-cased-finetuned-conll03-english",
-            device=0 if device == "cuda" else -1,
-            grouped_entities=True
+        classifier_pipeline = pipeline(
+            "zero-shot-classification",
+            model="valhalla/distilbart-mnli-12-3",
+            device=0 if device == "cuda" else -1
         )
-        logger.info("NER model loaded successfully!")
+        logger.info("Zero-shot classifier loaded")
     except Exception as e:
-        logger.error(f"Error loading NER model: {e}")
+        logger.error(f"Classifier load error: {e}")
 
-# Function to load the local language model
-def load_model():
-    """Load the local language model"""
+def load_qa_model():
+    """Load question-answering model on demand"""
     global qa_pipeline
-
     try:
-
-        # Check if GPU is available
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device: {device}")
-        
         qa_pipeline = pipeline(
             "question-answering",
             model="distilbert-base-uncased-distilled-squad",
             device=0 if device == "cuda" else -1
         )
-
-        logger.info("Model loaded successfully!")
-
+        logger.info("QA model loaded")
     except Exception as e:
-        logger.error(f"Error loading model: {e}")
+        logger.error(f"QA model load error: {e}")
 
-def search_web(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
-    """Search the web using DuckDuckGo"""
+def load_summarization_model():
+    """Load summarization model on demand"""
+    global summarization_pipeline
     try:
-        ddgs = DDGS()
-        results = []
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        summarization_pipeline = pipeline(
+            "summarization",
+            model="sshleifer/distilbart-cnn-6-6",
+            device=0 if device == "cuda" else -1
+        )
+        logger.info("Summarization model loaded")
+    except Exception as e:
+        logger.error(f"Summarization model load error: {e}")
 
-        for result in ddgs.text(query,safesearch='off',max_results=10):
-            results.append({
-                "title": result.get("title", ""),
-                "body": result.get("body", ""),
-                "href": result.get("href", "")
-            })
+def load_ner_model():
+    """Load NER model"""
+    global ner_pipeline
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        ner_pipeline = pipeline(
+            "ner",
+            model="dbmdz/bert-large-cased-finetuned-conll03-english",
+            device=0 if device == "cuda" else -1,
+            grouped_entities=True
+        )
+        logger.info("NER model loaded")
+    except Exception as e:
+        logger.error(f"NER model load error: {e}")
 
-        return results
+def classify_query(prompt: str) -> str:
+    """Classify query using zero-shot learning"""
+    candidate_labels = ['question answering', 'summarization']
+    try:
+        result = classifier_pipeline(
+            prompt,
+            candidate_labels,
+            multi_label=False
+        )
+        return result['labels'][0]
+    except Exception as e:
+        logger.error(f"Classification failed: {e}")
+        return 'question answering'  # Fallback to QA
 
+def search_web(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
+    """Enhanced web search with error handling"""
+    try:
+        with DDGS() as ddgs:
+            return [{
+                "title": r.get("title", ""),
+                "body": re.sub(r'\s+', ' ', r.get("body", "")).strip(),
+                "href": r.get("href", "")
+            } for r in ddgs.text(query, safesearch='off', max_results=max_results)]
     except Exception as e:
         logger.error(f"Search error: {e}")
         return []
 
-def generate_response(prompt: str,search_context: str) -> str:
-    """Generate response using local model"""
-    try:
-        if qa_pipeline is None:
-            return "Model not loaded properly. Please try again."
-
-        # Validate that qa_pipeline is a question-answering pipeline
-        if not hasattr(qa_pipeline, "task") or qa_pipeline.task != "question-answering":
-            return "Invalid pipeline type. Expected a question-answering pipeline."
-
-        result = qa_pipeline(question=prompt, context=search_context)['answer']
-        return result
-
-    except Exception as e:
-        logger.error(f"Generation error: {e}")
-        return f"Sorry, I encountered an error: {str(e)}"
+def format_search_context(results: List[Dict[str, Any]]) -> str:
+    """Create condensed context from search results"""
+    return "\n".join(
+        f"{i+1}. {res['title']}: {res['body'][:200]}"
+        for i, res in enumerate(results[:5])
+    )
 
 @app.on_event("startup")
 async def startup_event():
-    """Load model on startup"""
-    load_model()
+    """Initialize core models on startup"""
+    load_classifier()
     load_ner_model()
 
-@app.get("/")
-async def root():
-    """Health check endpoint"""
-    return {
-        "message": "Open Source Chat API is running!",
-        "model_loaded": qa_pipeline is not None,
-        "endpoints": {
-            "chat": "/chat",
-            "search": "/search",
-            "docs": "/docs"
-        }
-    }
-
 @app.post("/chat", response_model=ChatResponse)
-async def chat(request: ChatRequest):
-    """Main chat endpoint"""
+async def chat_endpoint(request: ChatRequest):
+    """Enhanced chat endpoint with dynamic model selection"""
     try:
-        search_results = None
-        search_context = None
+        search_results = []
+        search_context = ""
 
-        # Perform web search if requested
+        # Web search processing
         if request.use_search:
-
-            # Extract entities for focused search
             entities = ner_pipeline(request.prompt)
-            logger.info(f"Identified entities: {entities}")
-
-            # Create search query from entities
             search_terms = [
                 ent["word"] for ent in entities
                 if ent["entity_group"] in ["PER", "ORG", "LOC", "MISC"]
             ]
-            search_query = " ".join(search_terms) if search_terms else request.prompt
-            logger.info(f"Search query: {search_query}")
-
+            search_query = " ".join(search_terms) or request.prompt
             search_results = search_web(search_query)
-        
-            search_context = "\n".join([
-                    f"- {result['title']}: {result['body'][:200]}..."
-                    for result in search_results[:min(len(search_results), 5)]
-                ])
-             
+            search_context = format_search_context(search_results)
+
+        # Query classification
+        task_type = classify_query(request.prompt)
+        logger.info(f"Classified task: {task_type}")
+
+        # Dynamic model loading
+        if task_type == 'question answering':
+            if not qa_pipeline:
+                load_qa_model()
+            response = qa_pipeline(
+                question=request.prompt,
+                context=search_context or request.prompt,
+                max_answer_len=100
+            )['answer']
         else:
-            search_context = request.prompt
-
-        logger.info(f"Search context: {search_context}")
-        if search_context:
-        # Generate response
-            response = generate_response(
-                request.prompt,
-                search_context
+            if not summarization_pipeline:
+                load_summarization_model()
+            # Handle long contexts safely
+            inputs = summarization_pipeline.tokenizer(
+                search_context or request.prompt,
+                truncation=True,
+                max_length=1024,
+                return_tensors="pt"
             )
-        else:
-            response = "No context available to generate a response."
+            processed_context = summarization_pipeline.tokenizer.decode(
+                inputs['input_ids'][0],
+                skip_special_tokens=True
+            )
+            response = summarization_pipeline(
+                processed_context,
+                max_length=150,
+                min_length=30,
+                do_sample=False
+            )[0]['summary_text']
 
         return ChatResponse(
             response=response,
+            task_type=task_type,
             search_results=search_results if request.use_search else None
         )
 
     except Exception as e:
-        logger.error(f"Chat endpoint error: {e}")
+        logger.error(f"Chat error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/search")
-async def search(request: SearchRequest):
-    """Web search endpoint"""
+async def search_endpoint(request: SearchRequest):
+    """Search endpoint with improved error handling"""
     try:
-        results = search_web(request.query, request.max_results)
-        return {"results": results}
-
+        return {"results": search_web(request.query, request.max_results)}
     except Exception as e:
         logger.error(f"Search endpoint error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.get("/health")
+async def health_check():
+    """Enhanced health check with model status"""
+    return {
+        "status": "OK",
+        "models": {
+            "classifier": bool(classifier_pipeline),
+            "qa": bool(qa_pipeline),
+            "summarization": bool(summarization_pipeline),
+            "ner": bool(ner_pipeline)
+        }
+    }
+
 if __name__ == "__main__":
     import uvicorn
-    import os
-    host = os.getenv("HOST", "0.0.0.0")
-    port = int(os.getenv("PORT", 7860))  # Changed from 8000 to 7860
-    uvicorn.run(app, host=host, port=port)
+    uvicorn.run(app, host="0.0.0.0", port=7860)