import time
import tiktoken
import streamlit as st
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.base import Embeddings
from langchain.vectorstores import FAISS
from langchain.indexes import VectorstoreIndexCreator
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from typing import List
from together import Together
# from langchain.embeddings import TogetherEmbeddings
from langchain.schema import Document as LangchainDocument
st.set_page_config(page_title="چت بات ارتش", page_icon="🪖", layout="wide")
st.markdown("""
""", unsafe_allow_html=True)
col1, col2, col3 = st.columns([1, 1, 1])
with col2:
st.image("army.png", width=240)
st.markdown("""
""", unsafe_allow_html=True)
class TogetherEmbeddings(Embeddings):
def __init__(self, model_name: str, api_key: str):
self.model_name = model_name
self.client = Together(api_key=api_key)
def embed_documents(self, texts: List[str]) -> List[List[float]]:
response = self.client.embeddings.create(model=self.model_name, input=texts)
return [item.embedding for item in response.data]
def embed_query(self, text: str) -> List[float]:
return self.embed_documents([text])[0]
def count_tokens(text, model_name="gpt-3.5-turbo"):
enc = tiktoken.encoding_for_model(model_name)
return len(enc.encode(text))
@st.cache_resource
def get_pdf_index():
with st.spinner('📄 در حال پردازش فایل PDF...'):
loader = [PyPDFLoader('test1.pdf')]
pages = []
for l in loader:
pages.extend(l.load())
splitter_initial = RecursiveCharacterTextSplitter(
chunk_size=124,
chunk_overlap=25
)
small_chunks = []
for page in pages:
text = page.page_content
if len(text) > 124:
small_chunks.extend(splitter_initial.split_text(text))
else:
small_chunks.append(text)
final_chunks = []
max_tokens = 2000
for chunk in small_chunks:
token_count = count_tokens(chunk, model_name="gpt-3.5-turbo")
if token_count > max_tokens:
splitter_token_safe = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=100
)
smaller_chunks = splitter_token_safe.split_text(chunk)
final_chunks.extend(smaller_chunks)
else:
final_chunks.append(chunk)
documents = [LangchainDocument(page_content=text) for text in final_chunks]
embeddings = TogetherEmbeddings(
model_name="togethercomputer/m2-bert-80M-32k-retrieval",
api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
)
# اینجا دیگه Vectorstore مستقیم میسازیم با FAISS
vectordb = FAISS.from_documents(documents, embedding=embeddings)
return vectordb
index = get_pdf_index()
llm = ChatOpenAI(
base_url="https://api.together.xyz/v1",
api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979',
model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
)
chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type='stuff',
retriever=index.vectorstore.as_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
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("چطور میتونم کمک کنم؟")
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'question:پاسخ را فقط به زبان فارسی جواب بده {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