Update app.py

app.py

@@ -36,7 +36,7 @@ sorted(glob.glob('/content/anatomy_vol_*'))
 
 class CFG:
     # LLMs
-    model_name = 'llama2-13b-chat' # wizardlm, llama2-7b-chat, llama2-13b-chat, mistral-7B
+    model_name = 'llama2-13b-chat'  # wizardlm, llama2-7b-chat, llama2-13b-chat, mistral-7B
     temperature = 0
     top_p = 0.95
     repetition_penalty = 1.15
@@ -53,11 +53,12 @@ class CFG:
 
     # paths
     PDFs_path = '/content/'
-    Embeddings_path =  '/content/faiss-hp-sentence-transformers'
+    Embeddings_path = '/content/faiss-hp-sentence-transformers'
     Output_folder = './rag-vectordb'
 
-    def get_model(model = CFG.model_name):
 
+# Define model loading function
+def get_model(model=CFG.model_name):
     print('\nDownloading model: ', model, '\n\n')
 
     if model == 'wizardlm':
@@ -66,17 +67,17 @@ class CFG:
         tokenizer = AutoTokenizer.from_pretrained(model_repo)
 
         bnb_config = BitsAndBytesConfig(
-            load_in_4bit = True,
-            bnb_4bit_quant_type = "nf4",
-            bnb_4bit_compute_dtype = torch.float16,
-            bnb_4bit_use_double_quant = True,
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
         )
 
         model = AutoModelForCausalLM.from_pretrained(
             model_repo,
-            quantization_config = bnb_config,
-            device_map = 'auto',
-            low_cpu_mem_usage = True
+            quantization_config=bnb_config,
+            device_map='auto',
+            low_cpu_mem_usage=True
         )
 
         max_len = 1024
@@ -87,18 +88,18 @@ class CFG:
         tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True)
 
         bnb_config = BitsAndBytesConfig(
-            load_in_4bit = True,
-            bnb_4bit_quant_type = "nf4",
-            bnb_4bit_compute_dtype = torch.float16,
-            bnb_4bit_use_double_quant = True,
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
         )
 
         model = AutoModelForCausalLM.from_pretrained(
             model_repo,
-            quantization_config = bnb_config,
-            device_map = 'auto',
-            low_cpu_mem_usage = True,
-            trust_remote_code = True
+            quantization_config=bnb_config,
+            device_map='auto',
+            low_cpu_mem_usage=True,
+            trust_remote_code=True
         )
 
         max_len = 2048
@@ -109,23 +110,20 @@ class CFG:
         tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True)
 
         bnb_config = BitsAndBytesConfig(
-            load_in_4bit = True,
-            bnb_4bit_quant_type = "nf4",
-            bnb_4bit_compute_dtype = torch.float16,
-            bnb_4bit_use_double_quant = True,
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
         )
 
         model = AutoModelForCausalLM.from_pretrained(
             model_repo,
-            quantization_config = bnb_config,
-
-            low_cpu_mem_usage = True,
-            trust_remote_code = True
+            quantization_config=bnb_config,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True
         )
 
-        max_len = 2048  #8192
-        truncation=True,  # Explicitly enable truncation
-        padding="max_len"  # Optional: pad to max_length
+        max_len = 2048  # 8192
 
     elif model == 'mistral-7B':
         model_repo = 'mistralai/Mistral-7B-v0.1'
@@ -133,17 +131,17 @@ class CFG:
         tokenizer = AutoTokenizer.from_pretrained(model_repo)
 
         bnb_config = BitsAndBytesConfig(
-            load_in_4bit = True,
-            bnb_4bit_quant_type = "nf4",
-            bnb_4bit_compute_dtype = torch.float16,
-            bnb_4bit_use_double_quant = True,
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
         )
 
         model = AutoModelForCausalLM.from_pretrained(
             model_repo,
-            quantization_config = bnb_config,
-            device_map = 'auto',
-            low_cpu_mem_usage = True,
+            quantization_config=bnb_config,
+            device_map='auto',
+            low_cpu_mem_usage=True,
         )
 
         max_len = 1024
@@ -153,120 +151,26 @@ class CFG:
 
     return tokenizer, model, max_len
 
-def get_model(model = CFG.model_name):
-
-    print('\nDownloading model: ', model, '\n\n')
-
-    if model == 'wizardlm':
-        model_repo = 'TheBloke/wizardLM-7B-HF'
-
-        tokenizer = AutoTokenizer.from_pretrained(model_repo)
-
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit = True,
-            bnb_4bit_quant_type = "nf4",
-            bnb_4bit_compute_dtype = torch.float16,
-            bnb_4bit_use_double_quant = True,
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_repo,
-            quantization_config = bnb_config,
-            device_map = 'auto',
-            low_cpu_mem_usage = True
-        )
-
-        max_len = 1024
-
-    elif model == 'llama2-7b-chat':
-        model_repo = 'daryl149/llama-2-7b-chat-hf'
-
-        tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True)
-
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit = True,
-            bnb_4bit_quant_type = "nf4",
-            bnb_4bit_compute_dtype = torch.float16,
-            bnb_4bit_use_double_quant = True,
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_repo,
-            quantization_config = bnb_config,
-            device_map = 'auto',
-            low_cpu_mem_usage = True,
-            trust_remote_code = True
-        )
 
-        max_len = 2048
-
-    elif model == 'llama2-13b-chat':
-        model_repo = 'daryl149/llama-2-13b-chat-hf'
-
-        tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True)
-
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit = True,
-            bnb_4bit_quant_type = "nf4",
-            bnb_4bit_compute_dtype = torch.float16,
-            bnb_4bit_use_double_quant = True,
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_repo,
-            quantization_config = bnb_config,
-
-            low_cpu_mem_usage = True,
-            trust_remote_code = True
-        )
-
-        max_len = 2048  #8192
-        truncation=True,  # Explicitly enable truncation
-        padding="max_len"  # Optional: pad to max_length
-
-    elif model == 'mistral-7B':
-        model_repo = 'mistralai/Mistral-7B-v0.1'
-
-        tokenizer = AutoTokenizer.from_pretrained(model_repo)
-
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit = True,
-            bnb_4bit_quant_type = "nf4",
-            bnb_4bit_compute_dtype = torch.float16,
-            bnb_4bit_use_double_quant = True,
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_repo,
-            quantization_config = bnb_config,
-            device_map = 'auto',
-            low_cpu_mem_usage = True,
-        )
-
-        max_len = 1024
-
-    else:
-        print("Not implemented model (tokenizer and backbone)")
-
-    return tokenizer, model, max_len
-
-tokenizer, model, max_len = get_model(model = CFG.model_name)
+# Initialize model and tokenizer
+tokenizer, model, max_len = get_model(model=CFG.model_name)
 
+# Set up pipeline for LLM
 pipe = pipeline(
-    task = "text-generation",
-    model = model,
-    tokenizer = tokenizer,
-    pad_token_id = tokenizer.eos_token_id,
-#     do_sample = True,
-    max_length = max_len,
-    temperature = CFG.temperature,
-    top_p = CFG.top_p,
-    repetition_penalty = CFG.repetition_penalty
+    task="text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    pad_token_id=tokenizer.eos_token_id,
+    max_length=max_len,
+    temperature=CFG.temperature,
+    top_p=CFG.top_p,
+    repetition_penalty=CFG.repetition_penalty
 )
 
-### langchain pipeline
-llm = HuggingFacePipeline(pipeline = pipe)
+# Langchain pipeline
+llm = HuggingFacePipeline(pipeline=pipe)
 
+# Load PDFs from content
 loader = DirectoryLoader(
     CFG.PDFs_path,
     glob="./*.pdf",
@@ -277,13 +181,15 @@ loader = DirectoryLoader(
 
 documents = loader.load()
 
+# Split documents into chunks
 text_splitter = RecursiveCharacterTextSplitter(
-    chunk_size = CFG.split_chunk_size,
-    chunk_overlap = CFG.split_overlap
+    chunk_size=CFG.split_chunk_size,
+    chunk_overlap=CFG.split_overlap
 )
 
 texts = text_splitter.split_documents(documents)
 
+# Set up vector store with embeddings
 from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 
 vectordb = FAISS.from_documents(
@@ -291,10 +197,10 @@ vectordb = FAISS.from_documents(
     HuggingFaceEmbeddings(model_name='sentence-transformers/all-mpnet-base-v2')
 )
 
-### persist vector database
-vectordb.save_local(f"{CFG.Output_folder}/faiss_index_rag") # save in output folder
-#     vectordb.save_local(f"{CFG.Embeddings_path}/faiss_index_hp") # save in input folder
+# Persist vector database
+vectordb.save_local(f"{CFG.Output_folder}/faiss_index_rag")
 
+# Define prompt template for question answering
 prompt_template = """
 Don't try to make up an answer, if you don't know just say that you don't know.
 Answer in the same language the question was asked.
@@ -305,39 +211,35 @@ Use only the following pieces of context to answer the question at the end.
 Question: {question}
 Answer:"""
 
-
 PROMPT = PromptTemplate(
-    template = prompt_template,
-    input_variables = ["context", "question"]
+    template=prompt_template,
+    input_variables=["context", "question"]
 )
 
-retriever = vectordb.as_retriever(search_kwargs = {"k": CFG.k, "search_type" : "similarity"})
+# Set up retriever from vector store
+retriever = vectordb.as_retriever(search_kwargs={"k": CFG.k, "search_type": "similarity"})
 
+# Create the QA chain
 qa_chain = RetrievalQA.from_chain_type(
-    llm = llm,
-    chain_type = "stuff", # map_reduce, map_rerank, stuff, refine
-    retriever = retriever,
-    chain_type_kwargs = {"prompt": PROMPT},
-    return_source_documents = True,
-    verbose = False
+    llm=llm,
+    chain_type="stuff",  # map_reduce, map_rerank, stuff, refine
+    retriever=retriever,
+    chain_type_kwargs={"prompt": PROMPT},
+    return_source_documents=True,
+    verbose=False
 )
 
+# Function to wrap text to preserve newlines
 def wrap_text_preserve_newlines(text, width=700):
-    # Split the input text into lines based on newline characters
     lines = text.split('\n')
-
-    # Wrap each line individually
     wrapped_lines = [textwrap.fill(line, width=width) for line in lines]
-
-    # Join the wrapped lines back together using newline characters
     wrapped_text = '\n'.join(wrapped_lines)
-
     return wrapped_text
 
 
+# Function to process the LLM response
 def process_llm_response(llm_response):
     ans = wrap_text_preserve_newlines(llm_response['result'])
-
     sources_used = ' \n'.join(
         [
             source.metadata['source'].split('/')[-1][:-4]
@@ -346,10 +248,11 @@ def process_llm_response(llm_response):
             for source in llm_response['source_documents']
         ]
     )
-
     ans = ans + '\n\nSources: \n' + sources_used
     return ans
 
+
+# Function to get LLM response
 def llm_ans(query):
     start = time.time()
 
@@ -357,26 +260,28 @@ def llm_ans(query):
     ans = process_llm_response(llm_response)
 
     end = time.time()
-
     time_elapsed = int(round(end - start, 0))
     time_elapsed_str = f'\n\nTime elapsed: {time_elapsed} s'
     return ans + time_elapsed_str
 
+
+# Correct locale issue
 import locale
 locale.getpreferredencoding = lambda: "UTF-8"
 
+# Gradio interface
 import gradio as gr
 
- def predict(message, history):
-     # output = message # debug mode
+def predict(message, history):
+    output = str(llm_ans(message)).replace("\n", "<br/>")
+    return output
 
-     output = str(llm_ans(message)).replace("\n", "<br/>")
-     return output
-
- demo = gr.ChatInterface(
-     predict,
-     title = f' Open-Source LLM ({CFG.model_name}) Question Answering'
- )
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(label="Enter your question"),
+    outputs="html",
+    title=f'Open-Source LLM ({CFG.model_name}) Question Answering'
+)
 
- demo.queue()
- demo.launch()
+demo.queue()
+demo.launch()