app.py

new update

app.py

@@ -1,121 +1,121 @@
-# 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
-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
-
-# Download the model from HuggingFace
-model_name = "anakin87/zephyr-7b-alpha-sharded"
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.bfloat16,
-    quantization_config=bnb_config
-)
-
-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/'
-
-# Load the documents from Google Drive
-loader = PyPDFDirectoryLoader(folder_path)
-documents = loader.load()
-
-# 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")
-
-# Specify the retriever
-retriever = vectordb.as_retriever()
-
-# Build HuggingFace pipeline for using zephyr-7b-alpha
-pipeline = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    use_cache=True,
-    device_map="auto",
-    max_length=2048,
-    do_sample=True,
-    top_k=5,
-    num_return_sequences=1,
-    eos_token_id=tokenizer.eos_token_id,
-    pad_token_id=tokenizer.eos_token_id,
-)
-
-# Specify the llm
-llm = HuggingFacePipeline(pipeline=pipeline)
-
-# Define the create_conversation function
-def create_conversation(query: str, chat_history: list) -> tuple:
-    try:
-        memory = ConversationBufferMemory(
-            memory_key='chat_history',
-            return_messages=False
-        )
-        qa_chain = ConversationalRetrievalChain.from_llm(
-            llm=llm,
-            retriever=retriever,
-            memory=memory,
-            get_chat_history=lambda h: h,
-        )
-
-        result = qa_chain({'question': query, 'chat_history': chat_history})
-        chat_history.append((query, result['answer']))
-        return '', chat_history
-
-    except Exception as e:
-        chat_history.append((query, 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])
-
-# Launch the Gradio demo
-demo.launch()
+# 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
+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
+
+# Download the model from HuggingFace
+model_name = "anakin87/zephyr-7b-alpha-sharded"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    quantization_config=bnb_config
+)
+
+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/'
+
+# Load the documents from Google Drive
+loader = PyPDFDirectoryLoader(folder_path)
+documents = loader.load()
+
+# 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")
+
+# Specify the retriever
+retriever = vectordb.as_retriever()
+
+# Build HuggingFace pipeline for using zephyr-7b-alpha
+pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    use_cache=True,
+    device_map="auto",
+    max_length=2048,
+    do_sample=True,
+    top_k=5,
+    num_return_sequences=1,
+    eos_token_id=tokenizer.eos_token_id,
+    pad_token_id=tokenizer.eos_token_id,
+)
+
+# Specify the llm
+llm = HuggingFacePipeline(pipeline=pipeline)
+
+# Define the create_conversation function
+def create_conversation(query: str, chat_history: list) -> tuple:
+    try:
+        memory = ConversationBufferMemory(
+            memory_key='chat_history',
+            return_messages=False
+        )
+        qa_chain = ConversationalRetrievalChain.from_llm(
+            llm=llm,
+            retriever=retriever,
+            memory=memory,
+            get_chat_history=lambda h: h,
+        )
+
+        result = qa_chain({'question': query, 'chat_history': chat_history})
+        chat_history.append((query, result['answer']))
+        return '', chat_history
+
+    except Exception as e:
+        chat_history.append((query, 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])
+
+# Launch the Gradio demo
+demo.launch()