Update app.py

app.py

@@ -343,67 +343,66 @@ def answer_query(query_text, index, documents, llm_model, llm_tokenizer, embeddi
 
 
 
-if __name__ == '__main__':
-    import os
-    from faiss import write_index
-    import gradio as gr
-    import numpy as np
-    import torch 
-    from tqdm import tqdm
-    from torch.utils.data import DataLoader, Dataset
-    from datasets import load_dataset
-    import pandas as pd
-    import faiss
-    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoModel
-    from transformers import TextIteratorStreamer
-    from threading import Thread
+    # import os
+    # from faiss import write_index
+    # import gradio as gr
+    # import numpy as np
+    # import torch 
+    # from tqdm import tqdm
+    # from torch.utils.data import DataLoader, Dataset
+    # from datasets import load_dataset
+    # import pandas as pd
+    # import faiss
+    # from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoModel
+    # from transformers import TextIteratorStreamer
+    # from threading import Thread
     
-    torch.set_num_threads(2)
-    HF_TOKEN = os.environ.get("SECRET_TOKEN")  
+    # torch.set_num_threads(2)
+    # HF_TOKEN = os.environ.get("SECRET_TOKEN")  
     
-    class CFG:
-        embedding_model = 'TimKond/S-PubMedBert-MedQuAD'
-        batch_size = 128
-        device = ('cuda' if torch.cuda.is_available() else 'cpu')
-        llm = 'google/gemma-2b-it'
-        n_samples = 3
-
-    # Show config
-    config = CFG()
-    # config_items = {k: v for k, v in vars(CFG).items() if not k.startswith('__')}
-    # print(tabulate(config_items.items(), headers=['Parameter', 'Value'], tablefmt='fancy_grid'))
+class CFG:
+    embedding_model = 'TimKond/S-PubMedBert-MedQuAD'
+    batch_size = 128
+    device = ('cuda' if torch.cuda.is_available() else 'cpu')
+    llm = 'google/gemma-2b-it'
+    n_samples = 3
+
+# Show config
+config = CFG()
+# config_items = {k: v for k, v in vars(CFG).items() if not k.startswith('__')}
+# print(tabulate(config_items.items(), headers=['Parameter', 'Value'], tablefmt='fancy_grid'))
+
+
+# Obtener los datos y cargar o generar el índice
+df = get_all_data()
+documents = TextDataset(df)
+if not os.path.exists('./storage/faiss_index.faiss'):
+    embeddings = get_bert_embeddings(documents, CFG.batch_size, CFG.embedding_model, CFG.device)
+    index = create_faiss_index(embeddings)
+    write_index(index, './storage/faiss_index.faiss')
+else:
+    index = faiss.read_index('./storage/faiss_index.faiss')
+
+# Load the model
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
 
-    
-    # Obtener los datos y cargar o generar el índice
-    df = get_all_data()
-    documents = TextDataset(df)
-    if not os.path.exists('./storage/faiss_index.faiss'):
-        embeddings = get_bert_embeddings(documents, CFG.batch_size, CFG.embedding_model, CFG.device)
-        index = create_faiss_index(embeddings)
-        write_index(index, './storage/faiss_index.faiss')
-    else:
-        index = faiss.read_index('./storage/faiss_index.faiss')
-
-    # Load the model
-    quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.bfloat16
-    )
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it", token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", quantization_config=quantization_config, torch_dtype=torch.float16, low_cpu_mem_usage=True, token=HF_TOKEN)
 
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it", token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", quantization_config=quantization_config, torch_dtype=torch.float16, low_cpu_mem_usage=True, token=HF_TOKEN)
 
+def make_inference(query, hist):
+    return answer_query(query, index, documents, model, tokenizer, CFG.embedding_model, CFG.n_samples, CFG.device)
 
-    def make_inference(query, hist):
-        return answer_query(query, index, documents, model, tokenizer, CFG.embedding_model, CFG.n_samples, CFG.device)
-    
-    demo = gr.ChatInterface(fn = make_inference, 
-                    examples = ["What is diabetes?", "Is ginseng good for diabetes?", "What are the symptoms of diabetes?", "What is Celiac disease?"], 
-                    title = "Gemma 2b MedicalQA Chatbot", 
-                    description = "Gemma 2b Medical Chatbot is a chatbot that can help you with your medical queries. It is not a replacement for a doctor. Please consult a doctor for any medical advice.",
-                    )
-    demo.launch()
+demo = gr.ChatInterface(fn = make_inference, 
+                examples = ["What is diabetes?", "Is ginseng good for diabetes?", "What are the symptoms of diabetes?", "What is Celiac disease?"], 
+                title = "Gemma 2b MedicalQA Chatbot", 
+                description = "Gemma 2b Medical Chatbot is a chatbot that can help you with your medical queries. It is not a replacement for a doctor. Please consult a doctor for any medical advice.",
+                )
+demo.launch()