Update utils.py

utils.py

@@ -33,6 +33,29 @@ from langchain_experimental.pydantic_v1 import Extra, Field, root_validator
 
 emb_model = SentenceTransformer("all-MiniLM-L6-v2")
 
+class EmbeddingsSearch:
+    def __init__(self, metadata_df, emb_model):
+
+        self.model = emb_model
+        self.metadata_df = metadata_df
+        self.embeddings = self.model.encode(self.metadata_df['desc'].tolist())
+
+    def __call__(self, text: str, topk: int = 5):
+        q_emb = self.model.encode([text])
+        distances = cosine_similarity(q_emb, self.embeddings)
+        idx = np.flip(distances.argsort())[0]
+        distances.sort()
+        distances = np.flip(distances)[0]
+        
+        results = pd.DataFrame()
+        results['idx'] = idx.tolist()[:topk]
+        results['distances'] = distances.tolist()[:topk]
+
+        results['table'] = [
+            self.metadata_df.loc[i, "table"] for i in results['idx']
+        ]
+        return results
+        
 xls = pd.ExcelFile('SmartClever table explanations.xlsx')
 metadata_df = pd.DataFrame()
 i = 0
@@ -49,6 +72,8 @@ for sheet_name in xls.sheet_names:
     
     i += 1
 
+table_search = EmbeddingsSearch(metadata_df=metadata_df, emb_model=emb_model)
+
 def extract_question_type(llm, query):
     messages = [
     (
@@ -69,29 +94,6 @@ def extract_question_type(llm, query):
     else:
         return 'unknown'
         
-class EmbeddingsSearch:
-    def __init__(self, metadata_df, emb_model):
-
-        self.model = emb_model
-        self.metadata_df = metadata_df
-        self.embeddings = self.model.encode(self.metadata_df['desc'].tolist())
-
-    def __call__(self, text: str, topk: int = 5):
-        q_emb = self.model.encode([text])
-        distances = cosine_similarity(q_emb, self.embeddings)
-        idx = np.flip(distances.argsort())[0]
-        distances.sort()
-        distances = np.flip(distances)[0]
-        
-        results = pd.DataFrame()
-        results['idx'] = idx.tolist()[:topk]
-        results['distances'] = distances.tolist()[:topk]
-
-        results['table'] = [
-            self.metadata_df.loc[i, "table"] for i in results['idx']
-        ]
-        return results
-        
 warnings.filterwarnings('ignore', message="pandas only supports SQLAlchemy connectable.*", category=UserWarning, module='chain')
 
 intermediate_steps_KEY = "intermediate_steps"