From e1ac2511e25ac180285dc1809b6885a462774b8f Mon Sep 17 00:00:00 2001
From: Gautam <102940730+GautamR-Samagra@users.noreply.github.com>
Date: Tue, 19 Mar 2024 16:12:52 +0530
Subject: [PATCH] Update model.py

---
 src/embeddings/colbert/local/model.py | 33 ++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/embeddings/colbert/local/model.py b/src/embeddings/colbert/local/model.py
index 292d408..6c6105c 100644
--- a/src/embeddings/colbert/local/model.py
+++ b/src/embeddings/colbert/local/model.py
@@ -16,7 +16,38 @@ def __new__(cls, context):
         cls.df = pd.read_csv('/Testing1.csv')
         cls.df['PID'] = cls.df.index.astype(str)
         with Run().context(RunConfig(experiment='notebook')):
-            cls.searcher = Searcher(index='/colbert_agri_embeddings/', collection=cls.df['content'].to_list())
+            cls.searcher = Searcher(index='import pandas as pd
+from ragatouille import RAGPretrainedModel
+from request import ModelRequest
+from colbert import Indexer, Searcher
+from colbert.infra import Run, RunConfig, ColBERTConfig
+from colbert.data import Queries, Collection
+
+
+
+class Model():
+    def __new__(cls, context):
+        cls.context = context
+        if not hasattr(cls, 'instance'):
+            cls.instance = super(Model, cls).__new__(cls)
+        # Initialize Colbert
+        cls.df = pd.read_csv('Testing1.csv')
+        cls.df['PID'] = cls.df.index.astype(str)
+        with Run().context(RunConfig(experiment='notebook')):
+            cls.searcher = Searcher(index='/content/ai-tools/src/embeddings/colbert/local/colbert_akai/', collection=cls.df['content'].to_list())
+        print(cls.df.columns)
+        
+        return cls.instance
+
+    async def inference(self, request: ModelRequest):
+        query = request.text
+        k = request.k  
+        column_returned = 'id'  
+        results = self.searcher.search(query, k)
+        searched_ids = self.df.loc[results[0], column_returned].to_list()
+        searched_content = self.df.loc[results[0], 'content'].to_list()
+        return {"ids": searched_ids, "content": searched_content, "scores": results[2]}
+/', collection=cls.df['content'].to_list())
         print(cls.df.columns)
         
         return cls.instance