Update app.py

app.py

@@ -1,62 +1,87 @@
 import streamlit as st
-from langchain_community.document_loaders import PyPDFLoader
+from langchain.llms import HuggingFacePipeline
+from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
+from langchain.embeddings import SentenceTransformerEmbeddings
+from langchain.vectorstores import FAISS
 from langchain.chains import RetrievalQA
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
-from langchain_community.llms import HuggingFacePipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import os
 
-# Khởi tạo mô hình và tokenizer
-model_name = "google/flan-t5-base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-# Tạo pipeline cho HuggingFace
-pipe = pipeline(
-    "text2text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_length=512,
-    temperature=0,
-    repetition_penalty=1.15
-)
-
-llm = HuggingFacePipeline(pipeline=pipe)
-
-# Cấu hình Streamlit
-st.title("PDF Chatbot with Flan-T5")
-uploaded_file = st.file_uploader("Upload PDF", type="pdf")
+# Load FLAN-T5 model
+@st.cache_resource
+def load_llm():
+    model_name = "google/flan-t5-base"  # Adjust model size if needed
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    pipe = pipeline(
+        "text2text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_length=512,
+        temperature=0.7,  # Adjust for creativity
+        top_p=0.95,
+        repetition_penalty=1.15
+    )
+    return HuggingFacePipeline(pipeline=pipe)
 
-if uploaded_file:
-    # Lưu file tạm và load nội dung
-    with open("temp.pdf", "wb") as f:
-        f.write(uploaded_file.getbuffer())
-    
-    loader = PyPDFLoader("temp.pdf")
+# Process PDF and create vectorstore
+def process_pdf(pdf_path):
+    loader = PyPDFLoader(pdf_path)
     documents = loader.load()
     
-    # Chia nhỏ văn bản
-    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
+    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     texts = text_splitter.split_documents(documents)
     
-    # Tạo embeddings và vector store
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    db = FAISS.from_documents(texts, embeddings)
-    
-    # Tạo retrieval chain
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type="stuff",
-        retriever=db.as_retriever(search_kwargs={"k": 3}),
-        return_source_documents=True
-    )
+    embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+    vectorstore = FAISS.from_documents(texts, embeddings)
+    return vectorstore
+
+def main():
+    st.set_page_config(page_title="PDF Chatbot", page_icon="📄")
+    st.title("PDF Chatbot 📄")
+    st.markdown("Upload a PDF and ask questions about its content using FLAN-T5!")
+
+    uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
     
-    # Xử lý chat
-    question = st.text_input("Ask your question:")
-    if question:
-        result = qa_chain({"query": question})
-        st.write("Answer:", result["result"])
-        st.write("Sources:")
-        for doc in result['source_documents']:
-            st.write(f"- Page {doc.metadata['page']}: {doc.page_content[:200]}...")
+    if uploaded_file is not None:
+        # Save uploaded file temporarily
+        with open("temp.pdf", "wb") as f:
+            f.write(uploaded_file.getbuffer())
+        
+        # Process PDF
+        with st.spinner("Processing PDF..."):
+            vectorstore = process_pdf("temp.pdf")
+        
+        # Load LLM
+        llm = load_llm()
+        
+        # Create QA chain
+        qa_chain = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type="stuff",
+            retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
+            return_source_documents=True
+        )
+        
+        # Query input
+        query = st.text_input("Ask a question about the PDF:")
+        if query:
+            with st.spinner("Generating answer..."):
+                result = qa_chain({"query": query})
+                answer = result["result"]
+                source_docs = result["source_documents"]
+                
+                st.markdown("### Answer")
+                st.write(answer)
+                
+                with st.expander("Show Source Documents"):
+                    for i, doc in enumerate(source_docs):
+                        st.markdown(f"**Source {i+1}:**")
+                        st.write(doc.page_content)
+
+    else:
+        st.info("Please upload a PDF file to get started.")
+
+if __name__ == "__main__":
+    main()