created main.py

main.py

@@ -0,0 +1,119 @@
+from flask import Flask, request, jsonify
+import torch
+import numpy as np
+import pandas as pd
+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}")
+        
+        self.embedding_model = SentenceTransformer(
+            model_name_or_path="all-mpnet-base-v2",
+            device=self.device
+        )
+        
+        self.load_data(preprocessed_data_path)
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            "Qwen/Qwen2.5-1.5B-Instruct",
+            trust_remote_code=True
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            "Qwen/Qwen2.5-1.5B-Instruct",
+            device_map="auto",
+            trust_remote_code=True,
+            torch_dtype=torch.float16
+        )
+
+    def load_data(self, preprocessed_data_path: str):
+        df = pd.read_csv(preprocessed_data_path)
+        
+        def parse_embedding(embedding_str):
+            try:
+                embedding_str = embedding_str.strip()
+                if embedding_str.startswith('[') and embedding_str.endswith(']'):
+                    embedding_str = embedding_str[1:-1]
+                return np.fromstring(embedding_str, sep=',')
+            except:
+                print(f"Error parsing embedding: {embedding_str[:100]}...")
+                return None
+
+        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)
+
+    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)
+        
+        if len(query_embedding.shape) == 1:
+            query_embedding = query_embedding.unsqueeze(0)
+        
+        scores = util.dot_score(query_embedding, self.embeddings)[0]
+        _, indices = torch.topk(scores, k=min(n_chunks, len(self.chunks_data)))
+        indices = indices.cpu().numpy()
+        
+        return [
+            {
+                "sentence_chunk": self.chunks_data[i]["sentence_chunk"],
+                "Reference": self.chunks_data[i]["Reference"]
+            }
+            for i in indices
+        ]
+
+    def generate_response(self, query: str):
+        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
+            
+            with torch.no_grad():
+                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,
+                        repetition_penalty=1.05,
+                        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
+                    
+                    if outputs[0][-1] == self.tokenizer.eos_token_id:
+                        break
+            
+            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)})
+
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=5000)