app.py

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

# !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)
        
        return xc


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)[0]

PI = PineconeIndex()
PI.build_index()
qamodel = QAModel()
model, tokenizer = qamodel.load_sharded_model()

def request_answer(query):
    search_results = PI.search(query)
    answers = []
    # print(search_results)
    for r in search_results['matches']:
        if r['score'] >= 0.45:
            tokenized_context = tokenizer(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: """
            query_to_model = """You are doctor who knows cures to diseases. If you don't know, say you don't know. Please respond appropriately based on the context provided.\n\nContext: {}\n\n\nResponse: """
            for ind in range(0, len(tokenized_context['input_ids']), 512-42):                        
                decoded_tokens_for_context = tokenizer.batch_decode([tokenized_context['input_ids'][ind:ind+470]], skip_special_tokens=True)
                response = qamodel.query_model(model, tokenizer, query_to_model.format(decoded_tokens_for_context[0]))
                
                if not "don't know" in response:
                    answers.append(response)

    if len(answers) == 0:
        return 'Not enough information to answer the question'
    return '\n'.join(answers)


demo = gr.Interface(
    fn=request_answer,
    inputs=[
        gr.components.Textbox(label="User question(Response may take up to 2 mins because of hardware limitation)"),
    ],
    outputs=[
        gr.components.Textbox(label="Output (The answer is meant as a reference and not actual advice)"),
    ],
    cache_examples=True,
    title="MedQA assistant",
    #description="MedQA assistant"
)

demo.launch()