KhacHuy's picture
Update app.py
94826ad verified
# import os
# from PyPDF2 import PdfReader
# from langchain.text_splitter import RecursiveCharacterTextSplitter
# from langchain_google_genai import GoogleGenerativeAIEmbeddings
# import streamlit as st
# import google.generativeai as genai
# from langchain.vectorstores import FAISS
# from langchain.prompts import PromptTemplate
# from dotenv import load_dotenv
# from langchain_community.embeddings import SentenceTransformerEmbeddings
# from docx import Document # Thêm import để đọc file docx
# load_dotenv()
# genai.configure(api_key="AIzaSyC5hcS1goQ7emeXmyk_7eEQIie7J8OomC4") # Thay YOUR_API_KEY bằng API key của bạn
# model = genai.GenerativeModel('gemini-1.5-flash')
# # Đọc tất cả PDF và trả về văn bản
# def get_pdf_text(pdf_docs):
# text = ""
# for pdf in pdf_docs:
# pdf_reader = PdfReader(pdf)
# for page in pdf_reader.pages:
# text += page.extract_text() or ""
# return text
# # Đọc tất cả DOCX và trả về văn bản
# def get_docx_text(docx_docs):
# text = ""
# for doc in docx_docs:
# document = Document(doc)
# for paragraph in document.paragraphs:
# text += paragraph.text # Đảm bảo chuỗi này được đóng đúng cách
# return text
# # Tách văn bản thành các đoạn
# def get_text_chunks(text):
# splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
# chunks = splitter.split_text(text)
# return chunks
# # Tạo vector store từ các đoạn văn bản
# def get_vector_store(chunks):
# embeddings = SentenceTransformerEmbeddings(model_name="keepitreal/vietnamese-sbert", model_kwargs={"trust_remote_code": True})
# vector_store = FAISS.from_texts(chunks, embedding=embeddings)
# vector_store.save_local("faiss_index")
# # Tạo chuỗi hỏi đáp
# def create_qa_chain(prompt, db):
# def custom_llm(query, context):
# full_prompt = prompt.format(context=context, question=query)
# response = model.generate_content(full_prompt)
# if "câu trả lời không có trong ngữ cảnh" in response.text:
# response = model.generate_content(query)
# return response.text
# class CustomRetrievalQA:
# def __init__(self, retriever, prompt):
# self.retriever = retriever
# self.prompt = prompt
# def invoke(self, inputs):
# query = inputs["query"]
# docs = self.retriever.get_relevant_documents(query)
# context = " ".join([doc.page_content for doc in docs])
# answer = custom_llm(query, context)
# return {"answer": answer}
# retriever = db.as_retriever(search_kwargs={"k": 3}, max_tokens_limit=6000)
# return CustomRetrievalQA(retriever, prompt)
# def clear_chat_history():
# st.session_state.messages = [{"role": "assistant", "content": "Upload some PDFs or DOCs and ask me a question."}]
# def user_input(user_question):
# embeddings = SentenceTransformerEmbeddings(model_name="keepitreal/vietnamese-sbert", model_kwargs={"trust_remote_code": True})
# new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
# retriever = new_db.as_retriever()
# prompt_template = """
# Trả lời câu hỏi chi tiết nhất có thể từ ngữ cảnh được cung cấp. Nếu câu trả lời không nằm trong ngữ cảnh được cung cấp, hãy nói, "câu trả lời không có trong ngữ cảnh".
# Context:\n {context}\n
# Question: \n{question}\n
# Trả lời:
# """
# qa_chain = create_qa_chain(prompt_template, new_db)
# response = qa_chain.invoke({"query": user_question})
# return {"output_text": [response["answer"]]}
# def main():
# st.set_page_config(page_title="Gemini PDF & DOC Chatbot", page_icon="🤖")
# # Sidebar for uploading PDF and DOCX files
# with st.sidebar:
# st.title("Menu:")
# pdf_docs = st.file_uploader("Upload your PDF Files", type=["pdf"], accept_multiple_files=True)
# docx_docs = st.file_uploader("Upload your DOCX Files", type=["docx"], accept_multiple_files=True)
# if st.button("Submit & Process"):
# with st.spinner("Processing..."):
# raw_text = get_pdf_text(pdf_docs)
# raw_text += get_docx_text(docx_docs) # Kết hợp văn bản từ PDF và DOCX
# if raw_text:
# text_chunks = get_text_chunks(raw_text)
# get_vector_store(text_chunks)
# st.success("Done")
# else:
# st.error("No text extracted from the PDFs or DOCX files.")
# # Main content area for displaying chat messages
# st.title("Chat with PDF and DOCX files using Gemini🤖")
# st.write("Welcome to the chat!")
# st.sidebar.button('Clear Chat History', on_click=clear_chat_history)
# # Chat input
# if "messages" not in st.session_state.keys():
# st.session_state.messages = [{"role": "assistant", "content": "Upload some PDFs or DOCs and ask me a question."}]
# for message in st.session_state.messages:
# with st.chat_message(message["role"]):
# st.write(message["content"])
# if prompt := st.chat_input():
# st.session_state.messages.append({"role": "user", "content": prompt})
# with st.chat_message("user"):
# st.write(prompt)
# # Display chat messages and bot response
# if st.session_state.messages and st.session_state.messages[-1]["role"] != "assistant":
# with st.chat_message("assistant"):
# with st.spinner("Thinking..."):
# response = user_input(prompt)
# placeholder = st.empty()
# full_response = ''
# for item in response['output_text']:
# full_response += item
# placeholder.markdown(full_response)
# placeholder.markdown(full_response)
# if full_response:
# message = {"role": "assistant", "content": full_response}
# st.session_state.messages.append(message)
# if __name__ == "__main__":
# main()
import os
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import streamlit as st
import google.generativeai as genai
from langchain.vectorstores import FAISS
from langchain.prompts import PromptTemplate
from dotenv import load_dotenv
from langchain_community.embeddings import SentenceTransformerEmbeddings
from docx import Document # Thêm import để đọc file docx
load_dotenv()
genai.configure(api_key="AIzaSyC5hcS1goQ7emeXmyk_7eEQIie7J8OomC4") # Thay YOUR_API_KEY bằng API key của bạn
model = genai.GenerativeModel('gemini-1.5-flash')
# Đọc tất cả PDF và trả về văn bản
def get_pdf_text(pdf_docs):
text = ""
for pdf in pdf_docs:
pdf_reader = PdfReader(pdf)
for page in pdf_reader.pages:
text += page.extract_text() or ""
return text
# Đọc tất cả DOCX và trả về văn bản
def get_docx_text(docx_docs):
text = ""
for doc in docx_docs:
document = Document(doc)
for paragraph in document.paragraphs:
text += paragraph.text
return text
# Tách văn bản thành các đoạn
def get_text_chunks(text):
splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
chunks = splitter.split_text(text)
return chunks
# Tạo vector store từ các đoạn văn bản
def get_vector_store(chunks):
embeddings = SentenceTransformerEmbeddings(model_name="keepitreal/vietnamese-sbert", model_kwargs={"trust_remote_code": True})
vector_store = FAISS.from_texts(chunks, embedding=embeddings)
vector_store.save_local("faiss_index")
# Tạo chuỗi hỏi đáp
def create_qa_chain(prompt, db):
def custom_llm(query, context):
full_prompt = prompt.format(context=context, question=query)
response = model.generate_content(full_prompt)
if "Câu trả lời không có trong ngữ cảnh" in response.text:
response = model.generate_content(query)
return response.text
class CustomRetrievalQA:
def __init__(self, retriever, prompt):
self.retriever = retriever
self.prompt = prompt
def invoke(self, inputs):
query = inputs["query"]
docs = self.retriever.get_relevant_documents(query)
context = " ".join([doc.page_content for doc in docs])
answer = custom_llm(query, context)
return {"answer": answer}
retriever = db.as_retriever(search_kwargs={"k": 3}, max_tokens_limit=6000)
return CustomRetrievalQA(retriever, prompt)
def clear_chat_history():
st.session_state.messages = [{"role": "assistant", "content": "Upload some PDFs or DOCs and ask me a question."}]
def user_input(user_question):
embeddings = SentenceTransformerEmbeddings(model_name="keepitreal/vietnamese-sbert", model_kwargs={"trust_remote_code": True})
new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
retriever = new_db.as_retriever()
prompt_template = """
Trả lời câu hỏi chi tiết nhất có thể từ ngữ cảnh được cung cấp. Nếu câu trả lời không nằm trong ngữ cảnh được cung cấp, hãy nói, "Câu trả lời không có trong ngữ cảnh".
Context:\n {context}\n
Question: \n{question}\n
Trả lời:
"""
qa_chain = create_qa_chain(prompt_template, new_db)
response = qa_chain.invoke({"query": user_question})
return {"output_text": [response["answer"]]}
def main():
st.set_page_config(page_title="Gemini PDF & DOC Chatbot", page_icon="🤖")
# Sidebar for uploading PDF and DOCX files
with st.sidebar:
st.title("Menu:")
pdf_docs = st.file_uploader("Upload your PDF Files", type=["pdf"], accept_multiple_files=True)
docx_docs = st.file_uploader("Upload your DOCX Files", type=["docx"], accept_multiple_files=True)
if st.button("Submit & Process"):
with st.spinner("Processing..."):
raw_text = get_pdf_text(pdf_docs)
raw_text += get_docx_text(docx_docs) # Kết hợp văn bản từ PDF và DOCX
if raw_text:
text_chunks = get_text_chunks(raw_text)
get_vector_store(text_chunks)
st.success(f"Processed {len(pdf_docs)} PDFs and {len(docx_docs)} DOCs.")
else:
st.error("No text extracted from the PDFs or DOCX files.")
# Main content area for displaying chat messages
st.title("Chat with PDF and DOCX files using Gemini🤖")
st.write("Welcome to the chat!")
st.sidebar.button('Clear Chat History', on_click=clear_chat_history)
# Chat input
if "messages" not in st.session_state.keys():
st.session_state.messages = [{"role": "assistant", "content": "Upload some PDFs or DOCs and ask me a question."}]
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
if prompt := st.chat_input():
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.write(prompt)
# Display chat messages and bot response
if st.session_state.messages and st.session_state.messages[-1]["role"] != "assistant":
with st.chat_message("assistant"):
with st.spinner("Thinking..."):
response = user_input(prompt)
placeholder = st.empty()
full_response = ''
for item in response['output_text']:
full_response += item
placeholder.markdown(full_response)
placeholder.markdown(full_response)
if full_response:
message = {"role": "assistant", "content": full_response}
st.session_state.messages.append(message)
if __name__ == "__main__":
main()