Update app.py

app.py

@@ -8,136 +8,116 @@ import time
 
 import langchain
 
-### loaders
+# Loaders
 from langchain.document_loaders import PyPDFLoader, DirectoryLoader
 
-### splits
+# Splits
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
-### prompts
+# Prompts
 from langchain import PromptTemplate, LLMChain
 
-### vector stores
+# Vector stores
 from langchain.vectorstores import FAISS
 
-### models
+# Models
 from langchain.llms import HuggingFacePipeline
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 
-### retrievers
+# Retrievers
 from langchain.chains import RetrievalQA
 
 import torch
 import transformers
 from transformers import (
     AutoTokenizer, AutoModelForCausalLM,
-    BitsAndBytesConfig,
     pipeline
 )
 import gradio as gr
 import locale
-import time
-from langchain.embeddings.huggingface import HuggingFaceEmbeddings
-
-from transformers import logging
 import shutil
 
 # Clear transformers cache
-logging.set_verbosity_error()
+transformers.logging.set_verbosity_error()
 shutil.rmtree('./.cache', ignore_errors=True)
 
 class CFG:
-    # LLMs
-    model_name = 'llama2-13b-chat' # wizardlm, llama2-7b-chat, llama2-13b-chat, mistral-7B
+    # LLMs configuration
+    model_name = 'llama2-13b-chat'  # Options: wizardlm, llama2-7b-chat, llama2-13b-chat, mistral-7B
     temperature = 0
     top_p = 0.95
     repetition_penalty = 1.15
 
-    # splitting
+    # Text splitting configuration
     split_chunk_size = 800
     split_overlap = 0
 
-    # embeddings
+    # Embeddings configuration
     embeddings_model_repo = 'sentence-transformers/all-MiniLM-L6-v2'
 
-    # similar passages
+    # Similar passages configuration
     k = 6
 
-    # paths
+    # File paths configuration
     PDFs_path = './'
-    Embeddings_path =  './faiss-hp-sentence-transformers'
+    Embeddings_path = './faiss-hp-sentence-transformers'
     Output_folder = './rag-vectordb'
-    
+
 def get_model(model=CFG.model_name):
     print('\nDownloading model: ', model, '\n\n')
-    model_repo = None
-
-    if model == 'llama2-13b-chat':
-        model_repo = 'daryl149/llama-2-13b-chat-hf'
-
-    if model_repo:
-        tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_repo,
-            device_map="auto",
-            offload_folder="./offload",  # Specify offload folder
-            trust_remote_code=True
-        )
-        max_len = 2048
-    else:
+    
+    model_repo = 'daryl149/llama-2-13b-chat-hf' if model == 'llama2-13b-chat' else None
+    
+    if not model_repo:
         raise ValueError("Model not implemented: " + model)
 
+    tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_repo,
+        device_map="auto",
+        offload_folder="./offload",
+        trust_remote_code=True
+    )
+    
+    max_len = 2048
+    
     return tokenizer, model, max_len
 
-
 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
-
+    
+    return '\n'.join(wrapped_lines)
 
 def process_llm_response(llm_response):
     ans = wrap_text_preserve_newlines(llm_response['result'])
-
+    
     sources_used = ' \n'.join(
         [
-            source.metadata['source'].split('/')[-1][:-4]
-            + ' - page: '
-            + str(source.metadata['page'])
+            f"{source.metadata['source'].split('/')[-1][:-4]} - page: {source.metadata['page']}"
             for source in llm_response['source_documents']
         ]
     )
 
-    ans = ans + '\n\nSources: \n' + sources_used
-    return ans
+    return ans + '\n\nSources: \n' + sources_used
 
 def llm_ans(query):
     start = time.time()
-
+    
     llm_response = qa_chain.invoke(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'
+    
+    time_elapsed_str = f'\n\nTime elapsed: {int(round(end - start))} s'
+    
     return ans + time_elapsed_str
 
 def predict(message, history):
      output = str(llm_ans(message)).replace("\n", "<br/>")
      return output
 
-
-
-
-tokenizer, model, max_len = get_model(model = CFG.model_name)
+tokenizer, model, max_len = get_model(model=CFG.model_name)
 
 pipe = pipeline(
     task="text-generation",
@@ -150,46 +130,42 @@ pipe = pipeline(
     repetition_penalty=CFG.repetition_penalty
 )
 
-### langchain pipeline
-llm = HuggingFacePipeline(pipeline = pipe)
+# LangChain pipeline setup
+llm = HuggingFacePipeline(pipeline=pipe)
 
 loader = DirectoryLoader(
     CFG.PDFs_path,
     glob="./*.pdf",
     loader_cls=PyPDFLoader,
     show_progress=True,
-    use_multithreading=True
 )
 
 documents = loader.load()
 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)
 
 vectordb = FAISS.from_documents(
     texts,
-    HuggingFaceEmbeddings(model_name='sentence-transformers/all-mpnet-base-v2')
+    HuggingFaceInstructEmbeddings(model_name='sentence-transformers/all-mpnet-base-v2')
 )
 
-### persist vector database
+# Persist vector database
 vectordb.save_local(f"{CFG.Output_folder}/faiss_index_rag")
 
-retriever = vectordb.as_retriever(search_kwargs = {"k": CFG.k, "search_type" : "similarity"})
+retriever = vectordb.as_retriever(search_kwargs={"k": CFG.k})
 
 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",  # Options: map_reduce, map_rerank, stuff, refine
+    retriever=retriever,
 )
 
 prompt_template = """
-Don't try to make up an answer, if you don't know just say that you don't know.
+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.
 Use only the following pieces of context to answer the question at the end.
 
@@ -198,19 +174,17 @@ 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"]
 )
 
-
 locale.getpreferredencoding = lambda: "UTF-8"
 
 demo = gr.ChatInterface(
-     predict,
-     title = f' Open-Source LLM ({CFG.model_name}) Question Answering'
- )
+     fn=predict,
+     title=f'Open-Source LLM ({CFG.model_name}) Question Answering'
+)
 
 demo.queue()
-demo.launch()
+demo.launch()