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,12 +53,11 @@ 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':
@@ -67,17 +66,17 @@ def get_model(model=CFG.model_name):
         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
@@ -88,18 +87,18 @@ def get_model(model=CFG.model_name):
         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
@@ -110,20 +109,23 @@ def get_model(model=CFG.model_name):
         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
+        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'
@@ -131,17 +133,17 @@ def get_model(model=CFG.model_name):
         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
@@ -151,26 +153,120 @@ def get_model(model=CFG.model_name):
 
     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
+        )
 
-# Initialize model and tokenizer
-tokenizer, model, max_len = get_model(model=CFG.model_name)
+        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)
 
-# Set up pipeline for LLM
 pipe = pipeline(
-    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
+    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
 )
 
-# Langchain pipeline
-llm = HuggingFacePipeline(pipeline=pipe)
+### langchain pipeline
+llm = HuggingFacePipeline(pipeline = pipe)
 
-# Load PDFs from content
 loader = DirectoryLoader(
     CFG.PDFs_path,
     glob="./*.pdf",
@@ -181,15 +277,13 @@ 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(
@@ -197,10 +291,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")
+### 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
 
-# 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.
@@ -211,35 +305,39 @@ 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"]
 )
 
-# Set up retriever from vector store
-retriever = vectordb.as_retriever(search_kwargs={"k": CFG.k, "search_type": "similarity"})
+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]
@@ -248,11 +346,10 @@ 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()
 
@@ -260,28 +357,26 @@ 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 = str(llm_ans(message)).replace("\n", "<br/>")
-    return output
+ def predict(message, history):
+     # output = message # debug mode
 
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(label="Enter your question"),
-    outputs="html",
-    title=f'Open-Source LLM ({CFG.model_name}) Question Answering'
-)
+     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.queue()
-demo.launch()
+ demo.queue()
+ demo.launch() correct the code