Final Interface

app.py

@@ -1,51 +1,15 @@
 import time
 import streamlit as st
-from llama_index import ServiceContext, StorageContext, set_global_service_context, VectorStoreIndex
+from llama_index import ServiceContext, StorageContext, set_global_service_context, VectorStoreIndex, SimpleDirectoryReader, Document
 from llama_index.embeddings import LangchainEmbedding
 from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 from llama_index.llms import LlamaCPP
 from llama_index.llms.llama_utils import messages_to_prompt, completion_to_prompt
 from PyPDF2 import PdfReader
 
-# LLM Intialization
-llm = LlamaCPP(
-    model_url=None,  # We'll load locally.
-    # Trying small version of an already small model
-    model_path='./Models/phi-2.Q4_K_M.gguf',
-    temperature=0.1,
-    max_new_tokens=512,
-    context_window=2048,  # Phi-2 2K context window - this could be a limitation for RAG as it has to put the content into this context window
-    generate_kwargs={},
-    # set to at least 1 to use GPU
-    # This is small model and there's no indication of layers offloaded to the GPU
-    model_kwargs={"n_gpu_layers": 32},
-    messages_to_prompt=messages_to_prompt,
-    completion_to_prompt=completion_to_prompt,
-    verbose=True
-)
-
-# Embedding Initialization
-embed_model = LangchainEmbedding(
-    HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
-)
-
-# Service Context
-service_context = ServiceContext.from_defaults(
-    chunk_size=128,  # Number of tokens in each chunk
-    chunk_overlap=20,
-    # This should be automatically set with the model metadata but we'll force it to ensure wit is
-    context_window=2048,
-    num_output=768,  # Maximum output from the LLM, let's put this at 512 to ensure LlamaIndex saves that "space" for the output
-    llm=llm,
-    embed_model=embed_model
-)
-set_global_service_context(service_context)
-
-# Storage Context
-storage_context = StorageContext.from_defaults()
+# Model Specific Prompt
 
 
-# Model Specific Prompt
 def modelspecific_prompt(promptmessage):
     # Model Specific Prompt
     # As per https://huggingface.co/TheBloke/phi-2-GGUF
@@ -55,51 +19,101 @@ def modelspecific_prompt(promptmessage):
 # PDF to Text
 def extract_text_from_pdf(pdf):
     pdf_reader = PdfReader(pdf)
-    return ''.join(page.extract_text() for page in pdf_reader.pages)
-
-
-st.title("Llama-CPP Local LLM with RAG (Phi-2 RAG + TinyLlama CHAT)")
-
-
-pdf = st.file_uploader("Upload a PDF file", type=["pdf"])
-
-if pdf is not None:
-    documents = extract_text_from_pdf(pdf)
-    nodes = (service_context.node_parser.get_nodes_from_documents(documents))
-    storage_context.docstore.add_documents(nodes)
-    index = (VectorStoreIndex.from_documents(
-        documents, service_context=service_context, storage_context=storage_context, llm=llm))
-    chat_engine = index.as_chat_engine(chat_mode="simple", verbose=True)
-
-    # Initialize chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-
-    # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
-
-    # Accept user input
-    if prompt := st.chat_input("What is up?"):
-        # Add user message to chat history
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        # Display user message in chat message container
-        with st.chat_message("user"):
-            st.markdown(prompt)
-
-        # Display assistant response in chat message container
-        with st.chat_message("assistant"):
-            message_placeholder = st.empty()
-            full_response = ""
-            assistant_response = chat_engine.chat(modelspecific_prompt(prompt))
-            # Simulate stream of response with milliseconds delay
-            for chunk in assistant_response.split():
-                full_response += chunk + " "
-                time.sleep(0.05)
-                # Add a blinking cursor to simulate typing
-                message_placeholder.markdown(full_response + "▌")
-            message_placeholder.markdown(full_response)
-        # Add assistant response to chat history
-        st.session_state.messages.append(
-            {"role": "assistant", "content": full_response})
+    data = ''.join(page.extract_text() for page in pdf_reader.pages)
+    data = data.split()
+    return data
+
+
+def main():
+    # LLM Intialization
+    llm = LlamaCPP(
+        model_url=None,  # We'll load locally.
+        # Trying small version of an already small model
+        model_path='./Models/phi-2.Q4_K_M.gguf',
+        temperature=0.1,
+        max_new_tokens=512,
+        context_window=2048,  # Phi-2 2K context window - this could be a limitation for RAG as it has to put the content into this context window
+        generate_kwargs={},
+        # set to at least 1 to use GPU
+        # This is small model and there's no indication of layers offloaded to the GPU
+        model_kwargs={"n_gpu_layers": 32},
+        messages_to_prompt=messages_to_prompt,
+        completion_to_prompt=completion_to_prompt,
+        verbose=True
+    )
+
+    # Embedding Initialization
+    embed_model = LangchainEmbedding(
+        HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
+    )
+
+    # Service Context
+    service_context = ServiceContext.from_defaults(
+        chunk_size=128,  # Number of tokens in each chunk
+        chunk_overlap=20,
+        # This should be automatically set with the model metadata but we'll force it to ensure wit is
+        context_window=2048,
+        num_output=768,  # Maximum output from the LLM, let's put this at 512 to ensure LlamaIndex saves that "space" for the output
+        llm=llm,
+        embed_model=embed_model
+    )
+    set_global_service_context(service_context)
+
+    # Storage Context
+    storage_context = StorageContext.from_defaults()
+    st.title("Llama-CPP Local LLM with RAG (Phi-2 RAG)")
+
+    pdf = st.file_uploader("Upload a PDF file", type=["pdf"])
+
+    if pdf is not None:
+        text_list = extract_text_from_pdf(pdf)
+        documents = [Document(text=t) for t in text_list]
+        nodes = (service_context.node_parser.get_nodes_from_documents(documents))
+        storage_context.docstore.add_documents(nodes)
+        index = (VectorStoreIndex.from_documents(
+            documents, service_context=service_context, storage_context=storage_context, llm=llm))
+        chat_engine = index.as_chat_engine(chat_mode="simple", verbose=True)
+
+        # Initialize chat history
+        if "messages" not in st.session_state:
+            st.session_state.messages = []
+
+        # Display chat messages from history on app rerun
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+
+        # Accept user input
+        if prompt := st.chat_input("What is up?"):
+            # Add user message to chat history
+            st.session_state.messages.append(
+                {"role": "user", "content": prompt})
+            # Display user message in chat message container
+            with st.chat_message("user"):
+                st.markdown(prompt)
+
+            # Display assistant response in chat message container
+            with st.chat_message("assistant"):
+                message_placeholder = st.empty()
+                full_response = ""
+                assistant_response = chat_engine.chat(
+                    modelspecific_prompt(str(prompt)))
+                assistant_response = str(assistant_response)
+                # Simulate stream of response with milliseconds delay
+                for chunk in assistant_response.split():
+                    full_response += chunk + " "
+                    time.sleep(0.05)
+                    # Add a blinking cursor to simulate typing
+                    message_placeholder.markdown(full_response + "▌")
+                message_placeholder.markdown(full_response)
+            # Add assistant response to chat history
+            st.session_state.messages.append(
+                {"role": "assistant", "content": full_response})
+
+    # Credits
+    st.markdown(
+        "Made By Danyaal Majid & Muhammad Bin Asif [HuggingFace](https://huggingface.co/spaces/DanyaalMajid/NLP-Final-LocalLLM-RAG).")
+
+
+if __name__ == "__main__":
+    main()