Update app.py

app.py

@@ -3,18 +3,19 @@ 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_huggingface import HuggingFaceEmbeddings
-from langchain.chains import ConversationChain
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.memory import ConversationBufferMemory
 
 from pathlib import Path
 import chromadb
 from unidecode import unidecode
+
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import re
 
 # Constants
-LLM_MODEL = "t5-small"  # Changed to a Seq2Seq model compatible with AutoModelForSeq2SeqLM
+LLM_MODEL = "facebook/bart-large-cnn"  # Changed to a model with larger response capabilities
+LLM_MAX_TOKEN = 512
 DB_CHUNK_SIZE = 512
 CHUNK_OVERLAP = 24
 TEMPERATURE = 0.1
@@ -43,13 +44,12 @@ def create_db(splits, collection_name):
     return vectordb
 
 # Initialize langchain LLM chain
-def initialize_llmchain(llm_model, vector_db, progress=gr.Progress()):
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
     progress(0.5, desc="Initializing HF Hub...")
-
-    # Create the HuggingFacePipeline for the model
+    
     tokenizer = AutoTokenizer.from_pretrained(llm_model)
     model = AutoModelForSeq2SeqLM.from_pretrained(llm_model)
-    pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+    pipe = pipeline("summarization", model=model, tokenizer=tokenizer)
 
     progress(0.75, desc="Defining buffer memory...")
     memory = ConversationBufferMemory(
@@ -95,8 +95,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(vector_db, progress=gr.Progress()):
-    qa_chain = initialize_llmchain(LLM_MODEL, vector_db, progress)
+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)
     return qa_chain, "Complete!"
 
 def format_chat_history(message, chat_history):
@@ -165,7 +165,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(vector_db)
+            qa_chain, llm_status = initialize_LLM(TEMPERATURE, LLM_MAX_TOKEN, 20, 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])