|
import streamlit as st |
|
import time |
|
from langchain.document_loaders import PyPDFLoader |
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
from langchain_together import TogetherEmbeddings |
|
from langchain.vectorstores import FAISS |
|
from langchain.chat_models import ChatOpenAI |
|
from langchain.chains import RetrievalQA |
|
|
|
|
|
def batch_embed(texts, embeddings_model, batch_size=50): |
|
all_embeddings = [] |
|
for i in range(0, len(texts), batch_size): |
|
batch = texts[i:i+batch_size] |
|
embs = embeddings_model.embed_documents([doc.page_content for doc in batch]) |
|
all_embeddings.extend(embs) |
|
return all_embeddings |
|
|
|
@st.cache_resource |
|
def load_chunks_and_embeddings(): |
|
pdf_loader = PyPDFLoader('test1.pdf') |
|
pages = pdf_loader.load() |
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0) |
|
docs = text_splitter.split_documents(pages) |
|
|
|
embeddings = TogetherEmbeddings( |
|
api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979" |
|
) |
|
|
|
vectorstore = FAISS.from_documents([], embedding=embeddings) |
|
|
|
|
|
progress = st.progress(0, text="🔄 در حال پردازش چانکها...") |
|
total = len(docs) |
|
|
|
batch_size = 50 |
|
for i in range(0, total, batch_size): |
|
batch_docs = docs[i:i+batch_size] |
|
embeddings_batch = embeddings.embed_documents([doc.page_content for doc in batch_docs]) |
|
vectorstore.add_embeddings(embeddings_batch, batch_docs) |
|
|
|
progress.progress(min((i+batch_size)/total, 1.0)) |
|
|
|
progress.empty() |
|
return vectorstore |
|
|
|
|
|
with st.spinner("📚 در حال بارگذاری فایل و ساخت امبدینگها... لطفا صبور باشید"): |
|
vectorstore = load_chunks_and_embeddings() |
|
|
|
|
|
llm = ChatOpenAI( |
|
base_url="https://api.together.xyz/v1", |
|
api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979', |
|
model="meta-llama/Llama-3-70B-Instruct-Turbo-Free" |
|
) |
|
|
|
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 10}) |
|
|
|
chain = RetrievalQA.from_chain_type( |
|
llm=llm, |
|
chain_type='stuff', |
|
retriever=retriever, |
|
input_key='question' |
|
) |
|
|
|
|
|
if 'messages' not in st.session_state: |
|
st.session_state.messages = [] |
|
|
|
if 'pending_prompt' not in st.session_state: |
|
st.session_state.pending_prompt = None |
|
|
|
st.title("📄🤖 دستیار PDF شما") |
|
|
|
|
|
for msg in st.session_state.messages: |
|
with st.chat_message(msg['role']): |
|
st.markdown(f"🗨️ {msg['content']}", unsafe_allow_html=True) |
|
|
|
prompt = st.chat_input("سوالی از PDF داری؟") |
|
|
|
if prompt: |
|
st.session_state.messages.append({'role': 'user', 'content': prompt}) |
|
st.session_state.pending_prompt = prompt |
|
st.rerun() |
|
|
|
if st.session_state.pending_prompt: |
|
with st.chat_message('ai'): |
|
thinking = st.empty() |
|
thinking.markdown("🤖 در حال فکر کردن...") |
|
|
|
|
|
response = chain.run(f'فقط به زبان فارسی جواب بده. سوال: {st.session_state.pending_prompt}') |
|
answer = response.split("Helpful Answer:")[-1].strip() |
|
if not answer: |
|
answer = "متأسفم، اطلاعات دقیقی در این مورد ندارم." |
|
|
|
thinking.empty() |
|
full_response = "" |
|
placeholder = st.empty() |
|
|
|
for word in answer.split(): |
|
full_response += word + " " |
|
placeholder.markdown(full_response + "▌") |
|
time.sleep(0.03) |
|
|
|
placeholder.markdown(full_response) |
|
st.session_state.messages.append({'role': 'ai', 'content': full_response}) |
|
st.session_state.pending_prompt = None |
|
|