app.py

app.py

@@ -1,14 +1,8 @@
-# Install dependencies
-pip install -q transformers peft accelerate bitsandbytes safetensors sentencepiece streamlit chromadb langchain sentence-transformers gradio pypdf
-
 # Import necessary libraries
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
-
-import os
 import gradio as gr
 from google.colab import drive
-
 import chromadb
 from langchain.llms import HuggingFacePipeline
 from langchain.document_loaders import PyPDFDirectoryLoader
@@ -60,7 +54,7 @@ vectordb = Chroma.from_documents(documents=all_splits, embedding=embeddings, per
 retriever = vectordb.as_retriever()
 
 # Build HuggingFace pipeline for using zephyr-7b-alpha
-pipeline = pipeline(
+hf_pipeline = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
@@ -75,7 +69,7 @@ pipeline = pipeline(
 )
 
 # Specify the llm
-llm = HuggingFacePipeline(pipeline=pipeline)
+llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
 # Define the create_conversation function
 def create_conversation(query: str, chat_history: list) -> tuple:
@@ -96,26 +90,20 @@ def create_conversation(query: str, chat_history: list) -> tuple:
         return '', chat_history
 
     except Exception as e:
-        chat_history.append((query, e))
+        chat_history.append((query, str(e)))
         return '', chat_history
 
-def ask_question(query: str):
-    response = create_conversation(query, [])
-    gen_out = response[1][0][1]
-    response_start_token = "Helpful Answer:"
-    idx = gen_out.index(response_start_token)
-    rag_prompt = gen_out[:idx]
-    response_text = gen_out[idx:]
-
-    return rag_prompt, response_text
-
 # Define the Gradio UI
 with gr.Blocks() as demo:
     chatbot = gr.Chatbot(label='My Chatbot')
     msg = gr.Textbox()
     clear = gr.ClearButton([msg, chatbot])
 
-    msg.submit(create_conversation, [msg, chatbot], [msg, chatbot])
+    def submit_message(text):
+        _, chat_history = create_conversation(text, [])
+        chatbot.update(chat_history)
+
+    msg.submit(submit_message, [msg], [msg])
 
 # Launch the Gradio demo
 demo.launch()