army / app.py
M17idd's picture
Update app.py
d8bfdcf verified
raw
history blame
3.2 kB
import os
import concurrent.futures
import docx
import torch
import numpy as np
import streamlit as st
from hazm import *
from transformers import AutoTokenizer, AutoModel
# بارگذاری مدل
@st.cache_resource
def load_model():
tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-fa-base-uncased")
model = AutoModel.from_pretrained("HooshvareLab/bert-fa-base-uncased")
return tokenizer, model
tokenizer, model = load_model()
# پردازش فایل‌های Word و تبدیل به جملات
@st.cache_data
def load_text_chunks(folder_path):
normalizer = Normalizer()
sentence_tokenizer = SentenceTokenizer()
texts = []
for filename in os.listdir(folder_path):
if filename.endswith(".docx"):
full_path = os.path.join(folder_path, filename)
doc = docx.Document(full_path)
file_text = "\n".join([para.text for para in doc.paragraphs])
if file_text.strip():
texts.append(file_text)
all_sentences = []
for text in texts:
normalized = normalizer.normalize(text)
sentences = sentence_tokenizer.tokenize(normalized)
all_sentences.extend(sentences)
# تقسیم به بخش‌های ۵ جمله‌ای
chunks = []
for i in range(0, len(all_sentences), 5):
chunk = " ".join(all_sentences[i:i+5])
if chunk:
chunks.append(chunk)
return chunks
# محاسبه embedding با BERT
def get_embedding(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
embeddings = outputs.last_hidden_state.mean(dim=1)
return embeddings.squeeze().numpy()
# شباهت کسینوسی
def cosine_similarity(vec1, vec2):
return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
# رابط کاربری استریم‌لیت
st.title("🔎 یافتن نزدیک‌ترین بخش ۵ جمله‌ای به ورودی شما")
# مسیر پوشه فایل‌های docx
folder_path = '46'
# بارگذاری و نمایش تعداد بخش‌ها
chunks = load_text_chunks(folder_path)
st.success(f"{len(chunks)} بخش ۵ جمله‌ای بارگذاری شد.")
# ورودی کاربر
user_input = st.text_area("لطفاً جمله یا متن خود را وارد کنید:")
def calculate_similarities_parallel(user_embedding, chunks):
with concurrent.futures.ThreadPoolExecutor() as executor:
similarities = list(executor.map(lambda chunk: cosine_similarity(user_embedding, get_embedding(chunk)), chunks))
return similarities
if st.button("🔍 جستجو"):
if not user_input.strip():
st.warning("لطفاً یک جمله وارد کنید.")
else:
with st.spinner("در حال محاسبه شباهت‌ها..."):
user_embedding = get_embedding(user_input)
similarities = calculate_similarities_parallel(user_embedding, chunks)
most_similar_index = np.argmax(similarities)
result = chunks[most_similar_index]
st.subheader("📌 شبیه‌ترین بخش ۵ جمله‌ای:")
st.write(result)