Update app.py

app.py

@@ -1,87 +1,62 @@
 import streamlit as st
-from langchain.llms import HuggingFacePipeline
-from langchain.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import CharacterTextSplitter
-from langchain.embeddings import SentenceTransformerEmbeddings
-from langchain.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
 from langchain.chains import RetrievalQA
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-import os
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+from langchain_community.llms import HuggingFacePipeline
 
-# 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)
+# 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)
 
-# Process PDF and create vectorstore
-def process_pdf(pdf_path):
-    loader = PyPDFLoader(pdf_path)
+# Cấu hình Streamlit
+st.title("PDF Chatbot with Flan-T5")
+uploaded_file = st.file_uploader("Upload PDF", type="pdf")
+
+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")
     documents = loader.load()
     
-    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    # Chia nhỏ văn bản
+    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
     texts = text_splitter.split_documents(documents)
     
-    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")
+    # Tạo embeddings và vector store
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    db = FAISS.from_documents(texts, embeddings)
     
-    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()
+    # 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
+    )
+    
+    # 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]}...")