Update models/rag_system.py

models/rag_system.py

@@ -1,10 +1,9 @@
 import os
 import pandas as pd
-from transformers import pipeline
-from langchain_core.embeddings import HuggingFaceEmbeddings  # Updated import
-from langchain_community.vectorstores import FAISS
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.docstore.document import Document
+from transformers import pipeline, AutoTokenizer, AutoModel
+import torch
+import numpy as np
+from sentence_transformers import SentenceTransformer
 from utils.logger import setup_logger
 from utils.model_loader import ModelLoader
 
@@ -13,6 +12,7 @@ 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",
@@ -28,30 +28,34 @@ class RAGSystem:
             raise FileNotFoundError(f"CSV file not found at {csv_path}")
         
         try:
-            documents = pd.read_csv(csv_path)
-            docs = [
-                Document(
-                    page_content=str(row['Title']),
-                    metadata={'index': idx}
-                ) for idx, row in documents.iterrows()
-            ]
-            
-            text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-            split_docs = text_splitter.split_documents(docs)
-            
-            embeddings = HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-MiniLM-L6-v2"
+            self.documents = pd.read_csv(csv_path)
+            # Create embeddings for all documents
+            self.doc_embeddings = self.model.encode(
+                self.documents['Title'].astype(str).tolist(),
+                convert_to_tensor=True
             )
-            self.vector_store = FAISS.from_documents(split_docs, embeddings)
-            self.retriever = self.vector_store.as_retriever()
         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]
+
     def process_query(self, query):
         try:
-            retrieved_docs = self.retriever.get_relevant_documents(query)
-            retrieved_text = "\n".join([doc.page_content for doc in retrieved_docs])[:1000]
+            retrieved_docs = self.get_relevant_documents(query)
+            retrieved_text = "\n".join(retrieved_docs)[:1000]
             
             qa_input = {
                 "question": query,