Update rag_pipeline.py

rag_pipeline.py

@@ -1,65 +1,65 @@
-import faiss
-import numpy as np
-import pandas as pd
-from sentence_transformers import SentenceTransformer
-import re
-
-class QuoteRAG:
-    def __init__(self, model_path="models/fine_tuned_model", data_path="data/english_quotes.csv"):
-        # Load model
-        try:
-            self.model = SentenceTransformer(model_path)
-            print("Loaded fine-tuned model")
-        except:
-            self.model = SentenceTransformer("all-MiniLM-L6-v2")
-            print("Loaded base model")
-
-        # Load dataset
-        self.df = pd.read_csv(data_path)
-
-        # Encode all quotes
-        self.embeddings = self.model.encode(self.df["quote"].tolist(), convert_to_numpy=True)
-        d = self.embeddings.shape[1]
-
-        # Build FAISS index
-        self.index = faiss.IndexFlatL2(d)
-        self.index.add(self.embeddings)
-        print("FAISS index built with", len(self.df), "quotes")
-
-    def search(self, query, top_k=5):
-        # Encode query
-        query_emb = self.model.encode([query], convert_to_numpy=True)
-        distances, indices = self.index.search(query_emb, top_k * 3)  # fetch more for filtering
-
-        results = []
-        for idx, dist in zip(indices[0], distances[0]):
-            row = self.df.iloc[idx]
-
-            # Normalized similarity: 0–1 (higher is better)
-            similarity = round(1 / (1 + float(dist)), 3)
-
-            results.append({
-                "quote": row["quote"],
-                "author": row["author"],
-                "tags": row.get("tags", ""),
-                "similarity": similarity
-            })
-
-        # Simple author filter if author name is in query
-        query_lower = query.lower()
-        author_filtered = [r for r in results if r["author"].lower() in query_lower]
-
-        if author_filtered:
-            results = author_filtered[:top_k]
-        else:
-            results = results[:top_k]
-
-        return results
-
-
-if __name__ == "__main__":
-    rag = QuoteRAG()
-    query = "Quotes about insanity attributed to Einstein"
-    results = rag.search(query)
-    for r in results:
-        print(r)
+import faiss
+import numpy as np
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+import re
+
+class QuoteRAG:
+    def __init__(self, model_path="models/fine_tuned_model", data_path="english_quotes.csv"):
+        # Load model
+        try:
+            self.model = SentenceTransformer(model_path)
+            print("Loaded fine-tuned model")
+        except:
+            self.model = SentenceTransformer("all-MiniLM-L6-v2")
+            print("Loaded base model")
+
+        # Load dataset
+        self.df = pd.read_csv(data_path)
+
+        # Encode all quotes
+        self.embeddings = self.model.encode(self.df["quote"].tolist(), convert_to_numpy=True)
+        d = self.embeddings.shape[1]
+
+        # Build FAISS index
+        self.index = faiss.IndexFlatL2(d)
+        self.index.add(self.embeddings)
+        print("FAISS index built with", len(self.df), "quotes")
+
+    def search(self, query, top_k=5):
+        # Encode query
+        query_emb = self.model.encode([query], convert_to_numpy=True)
+        distances, indices = self.index.search(query_emb, top_k * 3)  # fetch more for filtering
+
+        results = []
+        for idx, dist in zip(indices[0], distances[0]):
+            row = self.df.iloc[idx]
+
+            # Normalized similarity: 0–1 (higher is better)
+            similarity = round(1 / (1 + float(dist)), 3)
+
+            results.append({
+                "quote": row["quote"],
+                "author": row["author"],
+                "tags": row.get("tags", ""),
+                "similarity": similarity
+            })
+
+        # Simple author filter if author name is in query
+        query_lower = query.lower()
+        author_filtered = [r for r in results if r["author"].lower() in query_lower]
+
+        if author_filtered:
+            results = author_filtered[:top_k]
+        else:
+            results = results[:top_k]
+
+        return results
+
+
+if __name__ == "__main__":
+    rag = QuoteRAG()
+    query = "Quotes about insanity attributed to Einstein"
+    results = rag.search(query)
+    for r in results:
+        print(r)