Update app.py

app.py

@@ -1,12 +1,9 @@
 import gradio as gr
-import os
-
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFacePipeline
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain_community.llms import HuggingFaceEndpoint
@@ -14,14 +11,11 @@ from langchain_community.llms import HuggingFaceEndpoint
 from pathlib import Path
 import chromadb
 from unidecode import unidecode
-
-from transformers import AutoTokenizer, AutoModel
-import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import re
 
-# Adjustments for the new LLM model
+# Constants
 LLM_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-LLM_MAX_TOKEN = 512
 DB_CHUNK_SIZE = 512
 CHUNK_OVERLAP = 24
 TEMPERATURE = 0.1
@@ -50,13 +44,13 @@ def create_db(splits, collection_name):
     return vectordb
 
 # Initialize langchain LLM chain
-def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+def initialize_llmchain(llm_model, vector_db, progress=gr.Progress()):
     progress(0.5, desc="Initializing HF Hub...")
 
-    # Use HuggingFacePipeline instead of HuggingFaceEndpoint
+    # Create the HuggingFacePipeline for the model
     tokenizer = AutoTokenizer.from_pretrained(llm_model)
-    model = AutoModel.from_pretrained(llm_model)
-    pipe = HuggingFacePipeline(model=model, tokenizer=tokenizer)
+    model = AutoModelForSeq2SeqLM.from_pretrained(llm_model)
+    pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
     progress(0.75, desc="Defining buffer memory...")
     memory = ConversationBufferMemory(
@@ -102,8 +96,8 @@ def initialize_database(pdf_url, chunk_size, chunk_overlap, progress=gr.Progress
     progress(0.9, desc="Done!")
     return vector_db, collection_name, "Complete!"
 
-def initialize_LLM(llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
-    qa_chain = initialize_llmchain(LLM_MODEL, llm_temperature, max_tokens, top_k, vector_db, progress)
+def initialize_LLM(vector_db, progress=gr.Progress()):
+    qa_chain = initialize_llmchain(LLM_MODEL, vector_db, progress)
     return qa_chain, "Complete!"
 
 def format_chat_history(message, chat_history):
@@ -172,7 +166,7 @@ def demo():
 
         def auto_initialize():
             vector_db, collection_name, db_status = initialize_database(pdf_url, DB_CHUNK_SIZE, CHUNK_OVERLAP)
-            qa_chain, llm_status = initialize_LLM(TEMPERATURE, LLM_MAX_TOKEN, 20, vector_db)
+            qa_chain, llm_status = initialize_LLM(vector_db)
             return vector_db, collection_name, db_status, qa_chain, llm_status, "Initialization complete."
 
         demo.load(auto_initialize, [], [vector_db, collection_name, db_progress, qa_chain, llm_progress])