Update main.py

main.py

@@ -1,17 +1,23 @@
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
 import torch
 import numpy as np
 import pandas as pd
 from sentence_transformers import SentenceTransformer, util
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from typing import Generator, List, Dict
 import json
+import asyncio
 
-app = Flask(__name__)
+app = FastAPI(title="Plant Chatbot API")
+
+class Query(BaseModel):
+    query: str
 
 class PlantChatbot:
     def __init__(self, preprocessed_data_path: str):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {self.device}")
         
         self.embedding_model = SentenceTransformer(
             model_name_or_path="all-mpnet-base-v2",
@@ -30,6 +36,9 @@ class PlantChatbot:
             trust_remote_code=True,
             torch_dtype=torch.float16
         )
+        
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
 
     def load_data(self, preprocessed_data_path: str):
         df = pd.read_csv(preprocessed_data_path)
@@ -46,13 +55,21 @@ class PlantChatbot:
 
         df["embedding"] = df["embedding"].apply(parse_embedding)
         df = df.dropna(subset=['embedding'])
+        
         self.chunks_data = df.to_dict(orient="records")
         
         embeddings_array = np.stack(df["embedding"].values)
-        self.embeddings = torch.tensor(embeddings_array, dtype=torch.float32).to(self.device)
+        self.embeddings = torch.tensor(
+            embeddings_array,
+            dtype=torch.float32
+        ).to(self.device)
 
-    def retrieve_relevant_chunks(self, query: str, n_chunks: int = 5) -> list[dict]:
-        query_embedding = self.embedding_model.encode(query, convert_to_tensor=True).to(self.device)
+    def retrieve_relevant_chunks(self, query: str, n_chunks: int = 5) -> List[Dict]:
+        query_embedding = self.embedding_model.encode(
+            query,
+            convert_to_tensor=True,
+            show_progress_bar=False
+        ).to(self.device)
         
         if len(query_embedding.shape) == 1:
             query_embedding = query_embedding.unsqueeze(0)
@@ -69,22 +86,41 @@ class PlantChatbot:
             for i in indices
         ]
 
-    def generate_response(self, query: str):
+    def format_prompt(self, query: str, context_chunks: List[Dict]) -> str:
+        context = "- " + "\n- ".join([chunk["sentence_chunk"] for chunk in context_chunks])
+        
+        prompt = f"""<|im_start|>system
+You are a knowledgeable plant expert. Provide helpful and accurate information about plants based on the given context.
+<|im_end|>
+<|im_start|>user
+Based on the following context about plants, please answer the query.
+Please be specific and detailed in your response, using only the information provided in the context.
+If you cannot answer the question based on the provided context, please say so.
+
+Context:
+{context}
+
+User query: {query}
+<|im_end|>
+<|im_start|>assistant
+"""
+        return prompt
+
+    async def generate_response(self, query: str) -> Generator[str, None, None]:
         try:
             context_chunks = self.retrieve_relevant_chunks(query)
             prompt = self.format_prompt(query, context_chunks)
-            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True)
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
-            
-            generated_text = ""
-            max_new_tokens = 4096
-            chunk_size = 50
+            input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
             
+            max_length = input_ids.shape[1] + 512
+
             with torch.no_grad():
-                for _ in range(0, max_new_tokens, chunk_size):
-                    outputs = self.model.generate(
-                        **inputs,
-                        max_new_tokens=chunk_size,
+                generated = input_ids
+                
+                while generated.shape[1] < max_length:
+                    outputs = self.model.forward(
+                        input_ids=generated,
+                        max_new_tokens=64,
                         do_sample=True,
                         temperature=0.7,
                         top_p=0.8,
@@ -92,28 +128,50 @@ class PlantChatbot:
                         pad_token_id=self.tokenizer.pad_token_id,
                         eos_token_id=self.tokenizer.eos_token_id,
                     )
-                    new_tokens = outputs[0][inputs['input_ids'].shape[1]:]
-                    new_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
-                    generated_text += new_text
-                    yield new_text
                     
-                    inputs['input_ids'] = outputs
+                    next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1).unsqueeze(0)
                     
-                    if outputs[0][-1] == self.tokenizer.eos_token_id:
+                    if next_token[0, 0].item() == self.tokenizer.eos_token_id:
                         break
+                        
+                    generated = torch.cat([generated, next_token.T], dim=1)
+                    new_text = self.tokenizer.decode(next_token[0], skip_special_tokens=True)
+                    
+                    if new_text:
+                        yield new_text
+                        await asyncio.sleep(0)  # Allow other tasks to run
             
             context_info = json.dumps({"context_items": context_chunks})
             yield f"\n<context>{context_info}</context>"
+            
         except Exception as e:
             print(f"Error generating response: {str(e)}")
-            yield "Error in processing your query."
+            yield "I apologize, but I encountered an error while processing your query. Please try again."
+
+# Initialize chatbot at startup
+chatbot = PlantChatbot("/home/avinashhn/bheri_bot/ravi/new/plant_data_chunks_and_embeddings.csv")
+
+async def response_generator(query: str):
+    """Wrapper generator for streaming response"""
+    async for chunk in chatbot.generate_response(query):
+        yield f"data: {chunk}\n\n"
+
+@app.post("/chat")
+async def chat(query: Query):
+    """Endpoint for chat interactions"""
+    try:
+        return StreamingResponse(
+            response_generator(query.query),
+            media_type="text/event-stream"
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-@app.route("/chat", methods=["POST"])
-def chat():
-    query = request.json.get('query')
-    chatbot = PlantChatbot("plant_data_chunks_and_embeddings.csv")
-    response = chatbot.generate_response(query)
-    return jsonify({"response": next(response)})
+@app.get("/")
+async def root():
+    """Root endpoint"""
+    return {"message": "Plant Chatbot API is running. Use /chat endpoint for queries."}
 
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=5000)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=5000)