Update main.py

main.py

@@ -1,4 +1,3 @@
-from flask import Flask, request, jsonify
 import torch
 import numpy as np
 import pandas as pd
@@ -6,20 +5,21 @@ from sentence_transformers import SentenceTransformer, util
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import json
 
-app = Flask(__name__)
-
 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}")
+        #print(f"Using device: {self.device}")
         
+        # Initialize embedding model
         self.embedding_model = SentenceTransformer(
             model_name_or_path="all-mpnet-base-v2",
             device=self.device
         )
         
+        # Load preprocessed data
         self.load_data(preprocessed_data_path)
         
+        # Initialize local Qwen model and tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(
             "Qwen/Qwen2.5-1.5B-Instruct",
             trust_remote_code=True
@@ -28,12 +28,18 @@ class PlantChatbot:
             "Qwen/Qwen2.5-1.5B-Instruct",
             device_map="auto",
             trust_remote_code=True,
-            torch_dtype=torch.float16
+            torch_dtype=torch.float16  # Use fp16 for memory efficiency
         )
+        
+        # Set pad token if not set
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
 
     def load_data(self, preprocessed_data_path: str):
+        """Load preprocessed data and prepare embeddings"""
         df = pd.read_csv(preprocessed_data_path)
         
+        # Convert embedding strings back to numpy arrays
         def parse_embedding(embedding_str):
             try:
                 embedding_str = embedding_str.strip()
@@ -46,13 +52,23 @@ class PlantChatbot:
 
         df["embedding"] = df["embedding"].apply(parse_embedding)
         df = df.dropna(subset=['embedding'])
+        
         self.chunks_data = df.to_dict(orient="records")
         
+        # Stack embeddings into a single tensor
         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)
+        """Retrieve relevant text chunks for the query"""
+        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 +85,51 @@ class PlantChatbot:
             for i in indices
         ]
 
+    def format_prompt(self, query: str, context_chunks: list[dict]) -> str:
+        """Format the prompt with context"""
+        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
+
     def generate_response(self, query: str):
+        """Generate a response for the user query"""
         try:
+            # Get relevant chunks
             context_chunks = self.retrieve_relevant_chunks(query)
+            
+            # Format prompt
             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
+            # Tokenize input
+            input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
+            
+            # Generate in smaller chunks
+            response = ""
+            max_length = input_ids.shape[1] + 512  # Limit total length
             
             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:
+                    # Generate next chunk
+                    outputs = self.model.forward(
+                        input_ids=generated,
+                        max_new_tokens=64,  # Generate smaller chunks at a time
                         do_sample=True,
                         temperature=0.7,
                         top_p=0.8,
@@ -92,28 +137,39 @@ 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
+                    # Get next token probabilities
+                    next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1).unsqueeze(0)
                     
-                    if outputs[0][-1] == self.tokenizer.eos_token_id:
+                    # Check for EOS token
+                    if next_token[0, 0].item() == self.tokenizer.eos_token_id:
                         break
+                        
+                    # Append new token and get its text
+                    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
             
+            # Add context information at the end
             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."
-
-@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)})
+            yield "I apologize, but I encountered an error while processing your query. Please try again."
 
+# Example usage
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=5000)
+    filepath = r"/home/avinashhn/bheri_bot/ravi/new/plant_data_chunks_and_embeddings.csv"
+    chatbot = PlantChatbot(filepath)
+    
+    # Example query
+    query = "Tell me about the Abidjan plant's care requirements"
+    
+    # Generate and print response
+    print("User:", query)
+    print("\nChatbot:")
+    for response_chunk in chatbot.generate_response(query):
+        print(response_chunk, end="", flush=True)