Update main.py

main.py

@@ -1,3 +1,4 @@
+from flask import Flask, request, jsonify
 import torch
 import numpy as np
 import pandas as pd
@@ -5,21 +6,20 @@ 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,18 +28,12 @@ class PlantChatbot:
             "Qwen/Qwen2.5-1.5B-Instruct",
             device_map="auto",
             trust_remote_code=True,
-            torch_dtype=torch.float16  # Use fp16 for memory efficiency
+            torch_dtype=torch.float16
         )
-        
-        # 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()
@@ -52,23 +46,13 @@ 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]:
-        """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)
+        query_embedding = self.embedding_model.encode(query, convert_to_tensor=True).to(self.device)
         
         if len(query_embedding.shape) == 1:
             query_embedding = query_embedding.unsqueeze(0)
@@ -85,51 +69,22 @@ 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()}
             
-            # 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
+            generated_text = ""
+            max_new_tokens = 4096
+            chunk_size = 50
             
             with torch.no_grad():
-                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
+                for _ in range(0, max_new_tokens, chunk_size):
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_new_tokens=chunk_size,
                         do_sample=True,
                         temperature=0.7,
                         top_p=0.8,
@@ -137,39 +92,28 @@ User query: {query}
                         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
                     
-                    # Get next token probabilities
-                    next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1).unsqueeze(0)
+                    inputs['input_ids'] = outputs
                     
-                    # Check for EOS token
-                    if next_token[0, 0].item() == self.tokenizer.eos_token_id:
+                    if outputs[0][-1] == 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 "I apologize, but I encountered an error while processing your query. Please try again."
+            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)})
 
-# Example usage
 if __name__ == "__main__":
-    filepath = r"C:\Users\neurotech\Desktop\NeuSpaarX\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)
+    app.run(host='0.0.0.0', port=5000)