The app

app.py

@@ -1,7 +1,203 @@
 import gradio as gr
+import pandas as pd
+import os
+from tqdm.auto import tqdm
+import pinecone
+from sentence_transformers import SentenceTransformer
+import torch
 
-def greet(name):
-    return "Hello " + name + "!!"
+from transformers import AutoModel, AutoConfig
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+# !pip install transformers accelerate
+# !pip install -qU pinecone-client[grpc] sentence-transformers
+# !pip install gradio
+
+class PineconeIndex:
+
+    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
+
+    def init_pinecone(self):
+
+        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'
+
+        pinecone.init(
+            api_key=PINECONE_API_KEY,
+            environment=PINECONE_ENV
+        )
+
+        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'
+            )
+
+        # now connect to the index
+        self.index = pinecone.GRPCIndex(index_name)
+        return self.index
+
+    def build_index(self):
+
+        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
+
+        sentence_model = self.sm
+
+        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 = []
+
+        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())
+
+            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]}")
+
+        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()
+
+    def search(self, query: str = "medicines for fever"):
+
+        sentence_model = self.sm
+
+        if self.index is None:
+            self.build_index()
+
+        index = self.index
+
+        # create the query vector
+        xq = sentence_model.encode(query).tolist()
+
+        # now query
+        xc = index.query(xq, top_k = 3, include_metadata = True)
+
+        results = []
+
+        for i in xc['matches']:
+            results.append(i['metadata']['text'])
+
+        return results
+
+
+
+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
+        
+        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()
+            
+        checkpoint = self.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)
+
+        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)
+
+PI = PineconeIndex()
+PI.build_index()
+qamodel = QAModel()
+model, tokenizer = qamodel.load_sharded_model()
+
+def request_answer(query):
+    search_results = PI.search(query)
+    answers = []
+    for r in search_results['matches']:
+        if r['score'] >= 0.45:
+            context = r['metadata']['text']
+            query_to_model = f"""You are doctor who knows cures to diseases. If you don't know the answer, please refrain from providing answers that are not relevant to the context. Please suggest appropriate remedies based on the context provided.\n\nContext: {context}\n\n\nResponse: """
+            answers.append(qamodel.query_model(model, tokenizer, query_to_model))
+
+    if len(answers) == 0:
+        return {'response': "Not enough information to answer the question"}['response']
+    return {'response': '\n'.join(answers)}['response']
+
+demo = gr.Interface(
+    fn=request_answer,
+    inputs=[
+        gr.components.Textbox(label="User question"),
+    ],
+    outputs=["text"],
+    examples=[["Building a demo with Gradio is so easy!",]],
+    cache_examples=False,
+    title="MedQA assistant",
+    #description="MedQA assistant"
+)
+
+if __name__ == "__main__":
+    print(demo.launch())