Update app.py

app.py

@@ -1,151 +1,96 @@
 import streamlit as st
-from langchain.llms import HuggingFacePipeline
+from streamlit_chat import message
+import tempfile
 from langchain.document_loaders import PyPDFLoader
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.embeddings import SentenceTransformerEmbeddings
+from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from langchain.chains import RetrievalQA
-from langchain.prompts import PromptTemplate
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
+from langchain.chains import ConversationalRetrievalChain
+from langchain_g4f import G4FLLM
+from g4f import Provider, models
 
-# --- 1. Load Mô Hình TinyLlama hoặc Mistral ---
-@st.cache_resource
-def load_llm():
-    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # Thay bằng "mistralai/Mistral-7B-Instruct-v0.2" nếu có GPU
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
-    if tokenizer.pad_token is None:
-        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+# Define the path for generated embeddings
+DB_FAISS_PATH = 'vectorstore/db_faiss'
 
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float32,  # Trên CPU nên dùng float32
-        device_map="auto"
-    )
-
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.9,
-        top_k=50,
-        repetition_penalty=1.2,
-        do_sample=True,
-        eos_token_id=tokenizer.eos_token_id,
-        truncation=True,
-        return_full_text=False
-    )
-
-    return HuggingFacePipeline(pipeline=pipe)
-
-
-# --- 2. Xử lý file PDF ---
-def process_pdf(pdf_path):
-    loader = PyPDFLoader(pdf_path)
-    documents = loader.load()
-    
-    text_splitter = CharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200,
-        separator="\n"
+# Load LLM using G4F (supports GPT-3.5, etc.)
+def load_llm():
+    llm = G4FLLM(
+        model=models.gpt_35_turbo,
+        provider=Provider.FreeGpt,
     )
-    texts = text_splitter.split_documents(documents)
-    
-    embeddings = SentenceTransformerEmbeddings(model_name="paraphrase-multilingual-mpnet-base-v2")
-    vectorstore = FAISS.from_documents(texts, embeddings)
-    return vectorstore
-
-
-# --- 3. Prompt Template tiếng Nhật (tự nhiên) ---
-template = """<s>[INST]あなたは親しみやすく丁寧なアシスタントです。以下の文書情報をもとに、質問に自然で分かりやすい日本語で回答してください。
-
-- 回答はできるだけ口語的で柔らかい表現を使ってください。
-- 理由や例を交えて説明すると良いでしょう。
-- 分からない場合は正直に「その点については詳しく記載されていません」と答えてください。
-
-文書情報:
-{context}
-
-質問: {question}
-回答: [/INST]"""
-
-QA_PROMPT = PromptTemplate(template=template, input_variables=["context", "question"])
-
-
-# --- 4. Hàm hậu xử lý câu trả lời ---
-def postprocess_answer(answer):
-    answer = answer.strip()
-    
-    for phrase in ["Answer:", "答え:", "回答:", "The answer is", "Based on the context"]:
-        answer = answer.replace(phrase, "").strip()
-
-    if answer and len(answer) > 0:
-        answer = answer[0].upper() + answer[1:]
-
-    if answer and answer[-1] not in "。.?！":
-        answer += "。"
-
-    if len(answer.split()) < 4:
-        answer = "資料にはその件についての詳細な記載が見受けられませんが、以下のように推測されます：" + answer
-
-    return answer
-
-
-# --- 5. Giao diện chính của ứng dụng ---
-def main():
-    st.set_page_config(page_title="PDFアシスタント", page_icon="📘")
-    st.title("PDFアシスタント 🤖")
-    st.markdown("PDFファイルをアップロードして内容について質問してください")
-
-    uploaded_file = st.file_uploader("PDFファイルを選択", type="pdf")
-    
-    if uploaded_file is not None:
-        with open("temp.pdf", "wb") as f:
-            f.write(uploaded_file.getbuffer())
-        
-        with st.spinner("ドキュメントを分析中..."):
-            vectorstore = process_pdf("temp.pdf")
-        
-        llm = load_llm()
-        
-        qa_chain = RetrievalQA.from_chain_type(
-            llm=llm,
-            chain_type="stuff",
-            retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
-            return_source_documents=True,
-            input_key="question",
-            chain_type_kwargs={
-                "prompt": QA_PROMPT,
-                "document_variable_name": "context"
-            }
-        )
-        
-        query = st.text_input("ドキュメントに関する質問を入力:")
-        if query:
-            with st.spinner("回答を生成中..."):
-                try:
-                    result = qa_chain({"question": query})
-                    raw_answer = result["result"]
-                    answer = postprocess_answer(raw_answer)
-                    
-                    st.markdown("### 回答")
-                    st.success(answer)
-                    
-                    with st.expander("参考資料を表示"):
-                        for i, doc in enumerate(result["source_documents"]):
-                            st.markdown(f"**引用 {i+1}:**")
-                            st.info(doc.page_content[:500] + "...")
-                            
-                except Exception as e:
-                    st.error(f"エラーが発生しました: {str(e)}")
-
-    else:
-        st.info("PDFファイルをアップロードしてください")
-
-
-if __name__ == "__main__":
-    main()
+    return llm
+
+# Hide default Streamlit style
+hide_streamlit_style = """
+        <style>
+        #MainMenu {visibility: hidden;}
+        footer {visibility: hidden;}
+        </style>
+        """
+st.markdown(hide_streamlit_style, unsafe_allow_html=True)
+
+# Set the title for the Streamlit app
+st.title("📄 PDF Chatbot - Zendo美女チャットボックス")
+
+# Upload PDF file
+uploaded_file = st.file_uploader("Tải lên tệp PDF của bạn", type="pdf")
+
+if uploaded_file is not None:
+    # Save temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmpfile:
+        tmpfile.write(uploaded_file.getvalue())
+        tmpfile_path = tmpfile.name
+
+    # Load PDF using PyPDFLoader
+    loader = PyPDFLoader(tmpfile_path)
+    pdf_data = loader.load()
+
+    # Create embeddings and save to FAISS
+    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2', model_kwargs={'device': 'cpu'})
+    db = FAISS.from_documents(pdf_data, embeddings)
+    db.save_local(DB_FAISS_PATH)
+
+    # Reload LLM and update retrieval chain
+    llm = load_llm()
+    chain = ConversationalRetrievalChain.from_llm(llm=llm, retriever=db.as_retriever())
+
+    st.success("PDF đã được tải lên và xử lý thành công!")
+
+    # Function for conversational chat
+    def conversational_chat(query):
+        result = chain({"question": query, "chat_history": st.session_state['history']})
+        st.session_state['history'].append((query, result["answer"]))
+        return result["answer"]
+
+    # Initialize session state for chat history and messages
+    if 'history' not in st.session_state:
+        st.session_state['history'] = []
+
+    if 'generated' not in st.session_state:
+        st.session_state['generated'] = ["こんにちは！zendo美女です。PDFの内容についてご質問ください... 🤗"]
+
+    if 'past' not in st.session_state:
+        st.session_state['past'] = ["チャットはここから"]
+
+    # UI Form for user input
+    response_container = st.container()
+    container = st.container()
+
+    with container:
+        with st.form(key='my_form', clear_on_submit=True):
+            user_input = st.text_input("ChatBox", placeholder="質問をご記入ください...", key='input')
+            submit_button = st.form_submit_button(label='Send')
+
+        if submit_button and user_input:
+            output = conversational_chat(user_input)
+            st.session_state['past'].append(user_input)
+            st.session_state['generated'].append(output)
+
+    # Display chat history
+    if st.session_state['generated']:
+        with response_container:
+            for i in range(len(st.session_state['generated'])):
+                message(st.session_state["past"][i], is_user=True, key=str(i) + '_user', avatar_style="big-smile")
+                message(st.session_state["generated"][i], key=str(i), avatar_style="thumbs")
+
+else:
+    st.info("Vui lòng tải lên một tệp PDF để bắt đầu trò chuyện.")