Update app.py

app.py

@@ -1,89 +1,98 @@
-
 import gradio as gr
 import os
 import tempfile
 import shutil
-import pdfminer.high_level
-import docx
-import faiss
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from sentence_transformers import SentenceTransformer
+from langchain_community.document_loaders import PyMuPDFLoader, UnstructuredWordDocumentLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
+from langchain_community.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+EMBEDDING_MODEL_NAME = "CAMeL-Lab/bert-base-arabic-camelbert-mix"
+QA_MODEL_NAME = "mosaicml/mpt-7b-storywriter"
 
-# تحميل النماذج
-embedding_model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2', device=device)
-qa_model_name = "aubmindlab/aragpt2-base"
-qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
-qa_model = AutoModelForCausalLM.from_pretrained(qa_model_name).to(device)
+embedding_model = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+qa_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_NAME, trust_remote_code=True)
+qa_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_NAME, trust_remote_code=True).to(device)
 
-# إعداد قاعدة البيانات
-index = None
-docs = []
+vectordb = None
 
-def extract_text(file_path):
-    if file_path.endswith('.pdf'):
-        with open(file_path, 'rb') as f:
-            return pdfminer.high_level.extract_text(f)
-    elif file_path.endswith('.docx') or file_path.endswith('.doc'):
-        doc = docx.Document(file_path)
-        return "\n".join([para.text for para in doc.paragraphs])
+def load_document(file_path):
+    ext = os.path.splitext(file_path)[1].lower()
+    if ext == ".pdf":
+        loader = PyMuPDFLoader(file_path)
+    elif ext in [".doc", ".docx"]:
+        loader = UnstructuredWordDocumentLoader(file_path)
     else:
-        raise ValueError("صيغة ملف غير مدعومة")
-
-def process_files(files):
-    global index, docs
-    all_text = ""
-    for file in files:
-        text = extract_text(file.name)
-        all_text += text + "\n"
+        raise ValueError("صيغة الملف غير مدعومة.")
+    return loader.load()
 
-    # تقسيم النص إلى مقاطع
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    texts = text_splitter.split_text(all_text)
+def train_from_documents(documents):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    texts = splitter.split_documents(documents)
+    vectorstore = FAISS.from_documents(texts, embedding_model)
+    return vectorstore
 
-    # إنشاء المتجهات
-    embeddings = embedding_model.encode(texts, show_progress_bar=True, convert_to_tensor=True)
-    index = faiss.IndexFlatL2(embeddings.shape[1])
-    index.add(embeddings.cpu().numpy())
-    docs = texts
-    return "✅ تم تحميل الكتب واستيعاب الأفكار! النظام جاهز للإجابة."
+def upload_files(files):
+    global vectordb
+    temp_dir = tempfile.mkdtemp()
+    all_docs = []
+    
+    for file in files:
+        file_path = os.path.join(temp_dir, file.name)
+        with open(file_path, "wb") as f:
+            f.write(file.read())
+        docs = load_document(file_path)
+        all_docs.extend(docs)
+    
+    vectordb = train_from_documents(all_docs)
+    shutil.rmtree(temp_dir)
+    return "✅ النظام جاهز للإجابة على أسئلتك!"
 
-def generate_answer(question):
-    global index, docs
-    if index is None:
-        return "❌ الرجاء رفع الكتب أولاً."
+def answer_question(question):
+    if vectordb is None:
+        return "⚠️ الرجاء رفع الملفات أولاً."
     
-    q_emb = embedding_model.encode([question])
-    D, I = index.search(q_emb, k=3)
-    context = "\n".join([docs[i] for i in I[0]])
+    retriever = vectordb.as_retriever(search_kwargs={"k": 5})
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=None,
+        retriever=retriever,
+        return_source_documents=True
+    )
 
-    # تجهيز الإدخال للنموذج
-    prompt = f"سؤال: {question}\n\nمحتوى ذو صلة:\n{context}\n\nالإجابة:"
-    inputs = qa_tokenizer(prompt, return_tensors='pt').to(device)
-    outputs = qa_model.generate(**inputs, max_new_tokens=300, pad_token_id=qa_tokenizer.eos_token_id)
+    relevant_docs = qa_chain.retriever.get_relevant_documents(question)
+    context = "\n".join(doc.page_content for doc in relevant_docs)
+    
+    inputs = qa_tokenizer(
+        f"أجب بالعربية فقط بناءً على السياق التالي:\n{context}\nالسؤال: {question}",
+        return_tensors="pt",
+        truncation=True,
+        max_length=1024
+    ).to(device)
+    
+    with torch.no_grad():
+        outputs = qa_model.generate(**inputs, max_length=300)
     answer = qa_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return answer.split("الإجابة:")[-1].strip()
+    return answer
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # 🚀 أهلاً بك في محاكاة عقل المؤلف
-    ارفع كتبك واستعد للانطلاق في رحلة استكشاف الأفكار العميقة!
-    """)
-    with gr.Tab("📚 رفع الكتب للتدريب"):
-        upload = gr.File(file_types=['.pdf', '.docx', '.doc'], file_count='multiple')
-        train_button = gr.Button("🚀 ابدأ التدريب!")
-        train_output = gr.Textbox(label="🔵 حالة التدريب", interactive=False)
-    with gr.Tab("❓ اسأل الكتاب"):
-        question = gr.Textbox(label="اكتب سؤالك هنا...")
-        answer = gr.Textbox(label="الإجابة", interactive=False)
-        ask_button = gr.Button("✉️ أرسل السؤال!")
+with gr.Blocks(title="محاكاة دماغ المؤلف") as demo:
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## 📚 ارفع كتبك هنا")
+            file_uploader = gr.File(file_types=[".pdf", ".doc", ".docx"], file_count="multiple")
+            upload_button = gr.Button("🚀 ابدأ التدريب")
+            training_status = gr.Textbox(label="حالة التدريب", interactive=False)
+        
+        with gr.Column():
+            gr.Markdown("## ❓ اطرح سؤالك")
+            question_input = gr.Textbox(label="سؤالك", placeholder="اكتب سؤالك هنا...")
+            ask_button = gr.Button("✉️ أرسل السؤال!")
+            answer_output = gr.Textbox(label="الإجابة", interactive=False)
 
-    train_button.click(process_files, inputs=[upload], outputs=[train_output])
-    ask_button.click(generate_answer, inputs=[question], outputs=[answer])
+    upload_button.click(upload_files, inputs=[file_uploader], outputs=[training_status])
+    ask_button.click(answer_question, inputs=[question_input], outputs=[answer_output])
 
-demo.launch()
+demo.launch(share=True)