Update app.py

app.py

@@ -5,7 +5,6 @@ from transformers import (
     pipeline,
     AutoModelForSequenceClassification,
     AutoTokenizer,
-    AutoModelForSeq2SeqLM,
     AutoModelForCausalLM,
     T5Tokenizer,
     T5ForConditionalGeneration,
@@ -128,7 +127,7 @@ class Chatbot:
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModelForCausalLM.from_pretrained(model_name)
 
-    def generate_response(self, prompt, max_length=50):
+    def generate_response(self, prompt, max_length=100):
         """
         Generates a response to a user query using GPT-2.
         """
@@ -137,6 +136,33 @@ class Chatbot:
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response
 
+    def handle_request(self, prompt):
+        """
+        Handles user requests by determining the intent and delegating to the appropriate function.
+        """
+        # Check if the user wants to search for something
+        if "search" in prompt.lower():
+            query = prompt.lower().replace("search", "").strip()
+            results = search_engine.search(query)
+            return {"type": "search", "results": results}
+
+        # Check if the user wants a summary
+        elif "summarize" in prompt.lower() or "summary" in prompt.lower():
+            text = prompt.lower().replace("summarize", "").replace("summary", "").strip()
+            summary = summarizer.summarize(text)
+            return {"type": "summary", "summary": summary}
+
+        # Check if the user wants to extract topics
+        elif "topics" in prompt.lower() or "topic" in prompt.lower():
+            text = prompt.lower().replace("topics", "").replace("topic", "").strip()
+            topics = topic_extractor.extract_topics([text])
+            return {"type": "topics", "topics": topics.to_dict()}
+
+        # Default to generating a conversational response
+        else:
+            response = self.generate_response(prompt)
+            return {"type": "chat", "response": response}
+
 
 # Initialize models
 classifier = ContentClassifier()
@@ -259,8 +285,18 @@ async def chat(request: PromptRequest):
     if not prompt:
         raise HTTPException(status_code=400, detail="No prompt provided")
 
-    response = chatbot.generate_response(prompt)
-    return {"response": response}
+    # Handle the request using the chatbot's handle_request method
+    result = chatbot.handle_request(prompt)
+
+    # Return the appropriate response based on the type of request
+    if result["type"] == "search":
+        return {"type": "search", "results": result["results"]}
+    elif result["type"] == "summary":
+        return {"type": "summary", "summary": result["summary"]}
+    elif result["type"] == "topics":
+        return {"type": "topics", "topics": result["topics"]}
+    else:
+        return {"type": "chat", "response": result["response"]}
 
 
 # Start the FastAPI app