remove colap

app.py

@@ -1,16 +1,14 @@
-# Import necessary libraries
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
 import gradio as gr
-from google.colab import drive
 import chromadb
-from langchain.llms import HuggingFacePipeline
 from langchain.document_loaders import PyPDFDirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
+from langchain_huggingface import HuggingFacePipeline
 
 # Download the model from HuggingFace
 model_name = "anakin87/zephyr-7b-alpha-sharded"
@@ -30,23 +28,18 @@ model = AutoModelForCausalLM.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 tokenizer.bos_token_id = 1  # Set beginning of sentence token id
 
-# Mount Google Drive and specify folder path
-drive.mount('/content/drive')
-folder_path = '/content/drive/MyDrive/TestcaseReport/'
+# Specify embedding model
+embedding_model_name = "sentence-transformers/all-mpnet-base-v2"
+model_kwargs = {"device": "cpu"}  # Using CPU since GPU is not available
+embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name, model_kwargs=model_kwargs)
 
-# Load the documents from Google Drive
-loader = PyPDFDirectoryLoader(folder_path)
-documents = loader.load()
+# Load the documents (replace this with your document loading logic)
+documents = ["Sample document text 1", "Sample document text 2"]
 
 # Split the documents into small chunks
 text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
 all_splits = text_splitter.split_documents(documents)
 
-# Specify embedding model
-embedding_model_name = "sentence-transformers/all-mpnet-base-v2"
-model_kwargs = {"device": "cpu"}  # Using CPU since GPU is not available
-embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name, model_kwargs=model_kwargs)
-
 # Embed document chunks
 vectordb = Chroma.from_documents(documents=all_splits, embedding=embeddings, persist_directory="chroma_db")
 
@@ -54,7 +47,7 @@ vectordb = Chroma.from_documents(documents=all_splits, embedding=embeddings, per
 retriever = vectordb.as_retriever()
 
 # Build HuggingFace pipeline for using zephyr-7b-alpha
-hf_pipeline = pipeline(
+pipeline = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
@@ -69,7 +62,7 @@ hf_pipeline = pipeline(
 )
 
 # Specify the llm
-llm = HuggingFacePipeline(pipeline=hf_pipeline)
+llm = HuggingFacePipeline(pipeline=pipeline)
 
 # Define the create_conversation function
 def create_conversation(query: str, chat_history: list) -> tuple:
@@ -90,7 +83,7 @@ def create_conversation(query: str, chat_history: list) -> tuple:
         return '', chat_history
 
     except Exception as e:
-        chat_history.append((query, str(e)))
+        chat_history.append((query, e))
         return '', chat_history
 
 # Define the Gradio UI
@@ -99,11 +92,7 @@ with gr.Blocks() as demo:
     msg = gr.Textbox()
     clear = gr.ClearButton([msg, chatbot])
 
-    def submit_message(text):
-        _, chat_history = create_conversation(text, [])
-        chatbot.update(chat_history)
-
-    msg.submit(submit_message, [msg], [msg])
+    msg.submit(create_conversation, [msg, chatbot], [msg, chatbot])
 
 # Launch the Gradio demo
 demo.launch()