Update app.py

app.py

@@ -32,6 +32,25 @@ QA_PROMPT = PromptTemplate(
   input_variables=["question", "context"]
 )
 
+# Load Phi-2 model from hugging face hub
+model_id = "microsoft/phi-2"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
+
+streamer = TextIteratorStreamer(tokenizer=tokenizer, skip_prompt=True)
+phi2 = pipeline(
+    "text-generation",
+    tokenizer=tokenizer,
+    model=model,
+    max_new_tokens=128,
+    pad_token_id=tokenizer.eos_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+    device_map="auto",
+    streamer=streamer
+  ) # GPU
+hf_model = HuggingFacePipeline(pipeline=phi2)
+
 # Returns a faiss vector store given a txt file
 def prepare_vector_store(filename):
   # Load data
@@ -54,42 +73,19 @@ def prepare_vector_store(filename):
 
   return vectorstore
 
-# Load Phi-2 model from hugging face hub
-model_id = "microsoft/phi-2"
-
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
-
-streamer = TextIteratorStreamer(tokenizer=tokenizer, skip_prompt=True)
-phi2 = pipeline(
-    "text-generation",
-    tokenizer=tokenizer,
-    model=model,
-    max_new_tokens=256,
-    eos_token_id=tokenizer.eos_token_id,
-    device_map="auto",
-    streamer=streamer
-  ) # GPU
-hf_model = HuggingFacePipeline(pipeline=phi2)
-
 # Retrieveal QA chian
-def get_retrieval_qa_chain(filename):
-  llm = hf_model
-  retriever = VectorStoreRetriever(
-    vectorstore=prepare_vector_store(filename)
-  )
-  model = RetrievalQA.from_chain_type(
-      llm=llm,
+def get_retrieval_qa_chain(retriever):
+  chain = RetrievalQA.from_chain_type(
+      llm=hf_model,
       retriever=retriever,
       chain_type_kwargs={"prompt": QA_PROMPT},
   )
-  return model
-
-# Question Answering Chain
-qa_chain = get_retrieval_qa_chain(filename="Oppenheimer-movie-wiki.txt")
+  return chain
 
 # Generates response using the question answering chain defined earlier
-def generate(question, answer):
+def generate(question, answer, retriever):
+  qa_chain = get_retrieval_qa_chain(retriever)
+
   query = f"{question}"
 
   thread = Thread(target=qa_chain.invoke, kwargs={"input": {"query": query}})
@@ -98,46 +94,56 @@ def generate(question, answer):
   response = ""
   for token in streamer:
     response += token
-    yield response
+    yield response.strip()
 
 # replaces the retreiver in the question answering chain whenever a new file is uploaded
-def upload_file(qa_chain):
-  def uploader(file):
-    qa_chain.retriever = VectorStoreRetriever(
+def upload_file(file):
+  new_retriever = VectorStoreRetriever(
       vectorstore=prepare_vector_store(file)
     )
-    return file
-  return uploader
+  return file, new_retriever
 
 with gr.Blocks() as demo:
   gr.Markdown("""
-  # RAG-Phi-2 Question Answering demo
-  ### This demo uses the Phi-2 language model and Retrieval Augmented Generation (RAG) to allow you to upload a txt file and ask the model questions related to the content of that file.
+  # Retrieval Augmented Generation with Phi-2: Question Answering demo
+  ### This demo uses the Phi-2 language model and Retrieval Augmented Generation (RAG). It allows you to upload a txt file and ask the model questions related to the content of that file.
   ### If you don't have one, there is a txt file already loaded, the new Oppenheimer movie's entire wikipedia page. The movie came out very recently in July, 2023, so the Phi-2 model is not aware of it.
   The context size of the Phi-2 model is 2048 tokens, so even this medium size wikipedia page (11.5k tokens) does not fit in the context window.
   Retrieval Augmented Generation (RAG) enables us to retrieve just the few small chunks of the document that are relevant to the our query and inject it into our prompt.
   The model is then able to answer questions by incorporating knowledge from the newly provided document. RAG can be used with thousands of documents, but this demo is limited to just one txt file.
   """)
 
-  file_output = gr.File()
+  default_text_file = "Oppenheimer-movie-wiki.txt"
+  retriever = gr.State(
+      VectorStoreRetriever(
+          vectorstore=prepare_vector_store(default_text_file)
+      )
+  )
+
+  gr.Markdown("## Upload a txt file or Use the Default 'Oppenheimer-movie-wiki.txt' that has already been loaded")
+
+  file_name = gr.Textbox(label="Loaded text file", value=default_text_file, lines=1, interactive=False)
   upload_button = gr.UploadButton(
       label="Click to upload a text file",
       file_types=["text"],
       file_count="single"
   )
-  upload_button.upload(upload_file(qa_chain), upload_button, file_output)
- 
+  upload_button.upload(upload_file, upload_button, [file_name, retriever])
+
+  gr.Markdown("## Enter your question")
+
   with gr.Row():
     with gr.Column():
       ques = gr.Textbox(label="Question", placeholder="Enter text here", lines=3)
     with gr.Column():
-      ans = gr.Textbox(label="Answer", lines=4)
+      ans = gr.Textbox(label="Answer", lines=4, interactive=False)
   with gr.Row():
     with gr.Column():
       btn = gr.Button("Submit")
     with gr.Column():
       clear = gr.ClearButton([ques, ans])
-  btn.click(fn=generate, inputs=[ques, ans], outputs=[ans])
+
+  btn.click(fn=generate, inputs=[ques, ans, retriever], outputs=[ans])
   examples = gr.Examples(
         examples=[
             "Who portrayed J. Robert Oppenheimer in the new Oppenheimer movie?",