army / app.py
M17idd's picture
Update app.py
d287817 verified
raw
history blame
3.95 kB
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
# --- 📄 ساخت امبدینگ‌ها با batch 50 تایی
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
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