Update models/rag_system.py

models/rag_system.py

@@ -1,9 +1,9 @@
 import os
-import pandas as pd
-from transformers import pipeline, AutoTokenizer, AutoModel
-import torch
 import numpy as np
+import pandas as pd
+from transformers import pipeline
 from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
 from utils.logger import setup_logger
 from utils.model_loader import ModelLoader
 
@@ -12,13 +12,18 @@ logger = setup_logger(__name__)
 class RAGSystem:
     def __init__(self, csv_path="apparel.csv"):
         try:
-            self.model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-            self.setup_system(csv_path)
-            self.qa_pipeline = ModelLoader.load_model_with_retry(
-                "distilbert-base-cased-distilled-squad",
-                pipeline,
-                task="question-answering"
+            # Initialize the sentence transformer model
+            self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+            
+            # Initialize the QA pipeline
+            self.qa_pipeline = pipeline(
+                "question-answering",
+                model="distilbert-base-cased-distilled-squad",
+                tokenizer="distilbert-base-cased-distilled-squad"
             )
+            
+            self.setup_system(csv_path)
+            
         except Exception as e:
             logger.error(f"Failed to initialize RAGSystem: {str(e)}")
             raise
@@ -28,41 +33,58 @@ class RAGSystem:
             raise FileNotFoundError(f"CSV file not found at {csv_path}")
         
         try:
+            # Load and preprocess documents
             self.documents = pd.read_csv(csv_path)
+            self.texts = self.documents['Title'].astype(str).tolist()
+            
             # Create embeddings for all documents
-            self.doc_embeddings = self.model.encode(
-                self.documents['Title'].astype(str).tolist(),
-                convert_to_tensor=True
-            )
+            self.embeddings = self.embedder.encode(self.texts)
+            
+            logger.info(f"Successfully loaded {len(self.texts)} documents")
+            
         except Exception as e:
             logger.error(f"Failed to setup RAG system: {str(e)}")
             raise
 
     def get_relevant_documents(self, query, top_k=5):
-        # Get query embedding
-        query_embedding = self.model.encode(query, convert_to_tensor=True)
-        
-        # Calculate cosine similarities
-        cos_scores = torch.nn.functional.cosine_similarity(
-            query_embedding.unsqueeze(0), 
-            self.doc_embeddings
-        )
-        
-        # Get top_k most similar documents
-        top_indices = torch.topk(cos_scores, min(top_k, len(self.documents))).indices
-        return [str(self.documents.iloc[idx]['Title']) for idx in top_indices]
+        try:
+            # Get query embedding
+            query_embedding = self.embedder.encode([query])
+            
+            # Calculate similarities
+            similarities = cosine_similarity(query_embedding, self.embeddings)[0]
+            
+            # Get top k most similar documents
+            top_indices = np.argsort(similarities)[-top_k:][::-1]
+            
+            return [self.texts[i] for i in top_indices]
+            
+        except Exception as e:
+            logger.error(f"Error retrieving relevant documents: {str(e)}")
+            return []
 
     def process_query(self, query):
         try:
-            retrieved_docs = self.get_relevant_documents(query)
-            retrieved_text = "\n".join(retrieved_docs)[:1000]
+            # Get relevant documents
+            relevant_docs = self.get_relevant_documents(query)
+            
+            if not relevant_docs:
+                return "No relevant documents found."
+            
+            # Combine retrieved documents into context
+            context = " ".join(relevant_docs)
             
+            # Prepare QA input
             qa_input = {
                 "question": query,
-                "context": retrieved_text
+                "context": context[:512]  # Limit context length for the model
             }
-            response = self.qa_pipeline(qa_input)
-            return response['answer']
+            
+            # Get answer using QA pipeline
+            answer = self.qa_pipeline(qa_input)
+            
+            return answer['answer']
+            
         except Exception as e:
-            logger.error(f"Query processing error: {str(e)}")
-            return "Failed to process query due to an error."
+            logger.error(f"Error processing query: {str(e)}")
+            return f"Failed to process query: {str(e)}"