Update app.py

app.py

@@ -1,167 +1,160 @@
 import gradio as gr
-import pandas as pd
 import os
-from tqdm.auto import tqdm
-import pinecone
-from sentence_transformers import SentenceTransformer
-import torch
-
-from transformers import AutoModel, AutoConfig
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+from pinecone_integration import PineconeIndex
+from qa_model import QAModel
 
 # !pip install transformers accelerate
 # !pip install -qU pinecone-client[grpc] sentence-transformers
 # !pip install gradio
 
-class PineconeIndex:
+# class PineconeIndex:
 
-    def __init__(self):
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+#     def __init__(self):
+#         device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-        self.sm = SentenceTransformer('all-MiniLM-L6-v2', device=device)
-        self.index_name = 'semantic-search-fast-med'
-        self.index = None
+#         self.sm = SentenceTransformer('all-MiniLM-L6-v2', device=device)
+#         self.index_name = 'semantic-search-fast-med'
+#         self.index = None
 
-    def init_pinecone(self):
+#     def init_pinecone(self):
 
-        index_name = self.index_name
-        sentence_model = self.sm
+#         index_name = self.index_name
+#         sentence_model = self.sm
 
-        # get api key from app.pinecone.io
-        PINECONE_API_KEY = "b97d5759-dd39-428b-a1fd-ed30f3ba74ee"  # os.environ.get('PINECONE_API_KEY') or 'PINECONE_API_KEY'
-        # find your environment next to the api key in pinecone console
-        PINECONE_ENV = "us-west4-gcp"  # os.environ.get('PINECONE_ENV') or 'PINECONE_ENV'
+#         # get api key from app.pinecone.io
+#         PINECONE_API_KEY = "b97d5759-dd39-428b-a1fd-ed30f3ba74ee"  # os.environ.get('PINECONE_API_KEY') or 'PINECONE_API_KEY'
+#         # find your environment next to the api key in pinecone console
+#         PINECONE_ENV = "us-west4-gcp"  # os.environ.get('PINECONE_ENV') or 'PINECONE_ENV'
 
-        pinecone.init(
-            api_key=PINECONE_API_KEY,
-            environment=PINECONE_ENV
-        )
+#         pinecone.init(
+#             api_key=PINECONE_API_KEY,
+#             environment=PINECONE_ENV
+#         )
 
-#         pinecone.delete_index(index_name)
+# #         pinecone.delete_index(index_name)
 
-        # only create index if it doesn't exist
-        if index_name not in pinecone.list_indexes():
-            pinecone.create_index(
-                name=index_name,
-                dimension=sentence_model.get_sentence_embedding_dimension(),
-                metric='cosine'
-            )
+#         # only create index if it doesn't exist
+#         if index_name not in pinecone.list_indexes():
+#             pinecone.create_index(
+#                 name=index_name,
+#                 dimension=sentence_model.get_sentence_embedding_dimension(),
+#                 metric='cosine'
+#             )
 
-        # now connect to the index
-        self.index = pinecone.GRPCIndex(index_name)
-        return self.index
+#         # now connect to the index
+#         self.index = pinecone.GRPCIndex(index_name)
+#         return self.index
 
-    def build_index(self):
+#     def build_index(self):
 
-        if self.index is None:
-            index = self.init_pinecone()
-        else:
-            index = self.index
+#         if self.index is None:
+#             index = self.init_pinecone()
+#         else:
+#             index = self.index
 
-        if index.describe_index_stats()['total_vector_count']:
-            "Index already built"
-            return
+#         if index.describe_index_stats()['total_vector_count']:
+#             "Index already built"
+#             return
 
-        sentence_model = self.sm
+#         sentence_model = self.sm
 
-        x = pd.read_excel('/kaggle/input/drug-p/Diseases_data_W.xlsx')
+#         x = pd.read_excel('/kaggle/input/drug-p/Diseases_data_W.xlsx')
 
-        question_dict = {'About': 'What is {}?', 'Symptoms': 'What are symptoms of {}?',
-                         'Causes': 'What are causes of {}?',
-                         'Diagnosis': 'What are diagnosis for {}?', 'Risk Factors': 'What are the risk factors for {}?',
-                         'Treatment Options': 'What are the treatment options for {}?',
-                         'Prognosis and Complications': 'What are the prognosis and complications?'}
-        context = []
-        disease_list = []
+#         question_dict = {'About': 'What is {}?', 'Symptoms': 'What are symptoms of {}?',
+#                          'Causes': 'What are causes of {}?',
+#                          'Diagnosis': 'What are diagnosis for {}?', 'Risk Factors': 'What are the risk factors for {}?',
+#                          'Treatment Options': 'What are the treatment options for {}?',
+#                          'Prognosis and Complications': 'What are the prognosis and complications?'}
+#         context = []
+#         disease_list = []
 
-        for i in range(len(x)):
-            disease = x.iloc[i, 0]
-            if disease.strip().lower() in disease_list:
-                continue
+#         for i in range(len(x)):
+#             disease = x.iloc[i, 0]
+#             if disease.strip().lower() in disease_list:
+#                 continue
 
-            disease_list.append(disease.strip().lower())
+#             disease_list.append(disease.strip().lower())
 
-            conditions = x.iloc[i, 1:].dropna().index
-            answers = x.iloc[i, 1:].dropna()
+#             conditions = x.iloc[i, 1:].dropna().index
+#             answers = x.iloc[i, 1:].dropna()
 
-            for cond in conditions:
-                context.append(f"{question_dict[cond].format(disease)}\n\n{answers[cond]}")
+#             for cond in conditions:
+#                 context.append(f"{question_dict[cond].format(disease)}\n\n{answers[cond]}")
 
-        batch_size = 128
-        for i in tqdm(range(0, len(context), batch_size)):
-            # find end of batch
-            i_end = min(i + batch_size, len(context))
-            # create IDs batch
-            ids = [str(x) for x in range(i, i_end)]
-            # create metadata batch
-            metadatas = [{'text': text} for text in context[i:i_end]]
-            # create embeddings
-            xc = sentence_model.encode(context[i:i_end])
-            # create records list for upsert
-            records = zip(ids, xc, metadatas)
-            # upsert to Pinecone
-            index.upsert(vectors=records)
+#         batch_size = 128
+#         for i in tqdm(range(0, len(context), batch_size)):
+#             # find end of batch
+#             i_end = min(i + batch_size, len(context))
+#             # create IDs batch
+#             ids = [str(x) for x in range(i, i_end)]
+#             # create metadata batch
+#             metadatas = [{'text': text} for text in context[i:i_end]]
+#             # create embeddings
+#             xc = sentence_model.encode(context[i:i_end])
+#             # create records list for upsert
+#             records = zip(ids, xc, metadatas)
+#             # upsert to Pinecone
+#             index.upsert(vectors=records)
 
-        # check number of records in the index
-        index.describe_index_stats()
+#         # check number of records in the index
+#         index.describe_index_stats()
 
-    def search(self, query: str = "medicines for fever"):
+#     def search(self, query: str = "medicines for fever"):
 
-        sentence_model = self.sm
+#         sentence_model = self.sm
 
-        if self.index is None:
-            self.build_index()
+#         if self.index is None:
+#             self.build_index()
 
-        index = self.index
+#         index = self.index
 
-        # create the query vector
-        xq = sentence_model.encode(query).tolist()
+#         # create the query vector
+#         xq = sentence_model.encode(query).tolist()
 
-        # now query
-        xc = index.query(xq, top_k = 3, include_metadata = True)
+#         # now query
+#         xc = index.query(xq, top_k = 3, include_metadata = True)
         
-        return xc
+#         return xc
 
 
-class QAModel():
-    def __init__(self, checkpoint="google/flan-t5-xl"):
-        self.checkpoint = checkpoint
-        self.tmpdir = f"{self.checkpoint.split('/')[-1]}-sharded"
+# class QAModel():
+#     def __init__(self, checkpoint="google/flan-t5-xl"):
+#         self.checkpoint = checkpoint
+#         self.tmpdir = f"{self.checkpoint.split('/')[-1]}-sharded"
 
-    def store_sharded_model(self):
-        tmpdir = self.tmpdir
+#     def store_sharded_model(self):
+#         tmpdir = self.tmpdir
         
-        checkpoint = self.checkpoint
+#         checkpoint = self.checkpoint
         
-        if not os.path.exists(tmpdir):
-            os.mkdir(tmpdir)
-            print(f"Directory created - {tmpdir}")
-            model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
-            print(f"Model loaded - {checkpoint}")
-            model.save_pretrained(tmpdir, max_shard_size="200MB")
-
-    def load_sharded_model(self):
-        tmpdir = self.tmpdir
-        if not os.path.exists(tmpdir):
-            self.store_sharded_model()
+#         if not os.path.exists(tmpdir):
+#             os.mkdir(tmpdir)
+#             print(f"Directory created - {tmpdir}")
+#             model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
+#             print(f"Model loaded - {checkpoint}")
+#             model.save_pretrained(tmpdir, max_shard_size="200MB")
+
+#     def load_sharded_model(self):
+#         tmpdir = self.tmpdir
+#         if not os.path.exists(tmpdir):
+#             self.store_sharded_model()
             
-        checkpoint = self.checkpoint
+#         checkpoint = self.checkpoint
         
 
-        config = AutoConfig.from_pretrained(checkpoint)
+#         config = AutoConfig.from_pretrained(checkpoint)
 
-        tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-        with init_empty_weights():
-            model = AutoModelForSeq2SeqLM.from_config(config)
-            # model = AutoModelForSeq2SeqLM.from_pretrained(tmpdir)
+#         tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+#         with init_empty_weights():
+#             model = AutoModelForSeq2SeqLM.from_config(config)
+#             # model = AutoModelForSeq2SeqLM.from_pretrained(tmpdir)
 
-        model = load_checkpoint_and_dispatch(model, checkpoint=tmpdir, device_map="auto")
-        return model, tokenizer
+#         model = load_checkpoint_and_dispatch(model, checkpoint=tmpdir, device_map="auto")
+#         return model, tokenizer
 
-    def query_model(self, model, tokenizer, query):
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        return tokenizer.batch_decode(model.generate(**tokenizer(query, return_tensors='pt').to(device)), skip_special_tokens=True)[0]
+#     def query_model(self, model, tokenizer, query):
+#         device = 'cuda' if torch.cuda.is_available() else 'cpu'
+#         return tokenizer.batch_decode(model.generate(**tokenizer(query, return_tensors='pt').to(device)), skip_special_tokens=True)[0]
 
 PI = PineconeIndex()
 PI.build_index()