Create rag_utils.py

rag_utils.py

@@ -0,0 +1,36 @@
+import os
+import faiss
+import torch
+from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
+from PyPDF2 import PdfReader
+
+class RAGRetriever:
+    def __init__(self):
+        self.encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+        self.index = faiss.IndexFlatL2(384)
+        self.contexts = []
+        self.ids = []
+
+    def add_document(self, text):
+        sentences = text.split("\n")
+        clean_sentences = [s.strip() for s in sentences if s.strip()]
+        embeddings = self.encoder.encode(clean_sentences)
+        self.index.add(embeddings)
+        self.contexts.extend(clean_sentences)
+
+    def retrieve(self, query, top_k=3):
+        q_vec = self.encoder.encode([query])
+        D, I = self.index.search(q_vec, top_k)
+        return [self.contexts[i] for i in I[0]]
+
+def extract_text_from_file(file_path):
+    ext = os.path.splitext(file_path)[-1].lower()
+    if ext == ".txt":
+        with open(file_path, "r", encoding="utf-8") as f:
+            return f.read()
+    elif ext == ".pdf":
+        reader = PdfReader(file_path)
+        return "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
+    else:
+        return ""