Update app.py

app.py

@@ -1,122 +1,117 @@
+# app.py
+
 import gradio as gr
 import os
-import tempfile
-import faiss
-import torch
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.prompts import PromptTemplate
-from langchain.chains import RetrievalQA
-from langchain.llms import HuggingFacePipeline
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from pdfminer.high_level import extract_text as extract_pdf_text
+import pdfminer.high_level
 import docx
-import nltk
-
-nltk.download('punkt')
-from nltk.tokenize import sent_tokenize
-
-uploaded_texts = []
-vector_store = None
-qa_chain = None
-
-embedding_model_name = "CAMeL-Lab/bert-base-arabic-camelbert-mix"
-embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)
-
-model_name = "csebuetnlp/mT5_small"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer, max_length=512)
-llm = HuggingFacePipeline(pipeline=pipe)
-
-ARABIC_PROMPT_TEMPLATE = """
-أنت نظام ذكي يجيب بناءً فقط على المعلومات المستخرجة من الكتب.
-لا تستخدم أي معلومات خارجية.
-السؤال: {question}
-الإجابة:
-"""
-
-def format_arabic_prompt(question):
-    return ARABIC_PROMPT_TEMPLATE.format(question=question)
-
-def extract_text_from_file(file_path):
-    if file_path.endswith(".pdf"):
-        return extract_pdf_text(file_path)
-    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])
-    else:
-        raise ValueError("Unsupported file format")
-
-def arabic_split_text(text):
-    sentences = sent_tokenize(text, language='arabic')
+from sentence_transformers import SentenceTransformer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import faiss
+import tempfile
+
+# ====== Settings ======
+EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+GENERATION_MODEL_NAME = "aubmindlab/aragpt2-small"
+CHUNK_SIZE = 500
+CHUNK_OVERLAP = 50
+TOP_K = 5
+
+# ====== Load Models ======
+embedder = SentenceTransformer(EMBEDDING_MODEL_NAME)
+gen_tokenizer = AutoTokenizer.from_pretrained(GENERATION_MODEL_NAME)
+gen_model = AutoModelForCausalLM.from_pretrained(GENERATION_MODEL_NAME)
+
+# ====== Globals ======
+index = None
+chunks = []
+
+# ====== Helpers ======
+def extract_text_from_pdf(file_path):
+    with open(file_path, 'rb') as f:
+        return pdfminer.high_level.extract_text(f)
+
+def extract_text_from_docx(file_path):
+    doc = docx.Document(file_path)
+    return "\n".join([para.text for para in doc.paragraphs])
+
+def chunk_text(text):
+    words = text.split()
     chunks = []
-    chunk = ""
-    for sentence in sentences:
-        if len(chunk) + len(sentence) <= 500:
-            chunk += " " + sentence
-        else:
-            chunks.append(chunk.strip())
-            chunk = sentence
-    if chunk:
-        chunks.append(chunk.strip())
+    for i in range(0, len(words), CHUNK_SIZE - CHUNK_OVERLAP):
+        chunk = " ".join(words[i:i+CHUNK_SIZE])
+        chunks.append(chunk)
     return chunks
 
-def train_from_texts(texts):
-    global vector_store, qa_chain
-
-    splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=100,
-        length_function=len,
-    )
-    
-    all_chunks = []
-    for text in texts:
-        chunks = arabic_split_text(text)
-        all_chunks.extend(chunks)
-
-    vectors = embeddings.embed_documents(all_chunks)
-    dimension = len(vectors[0])
-    index = faiss.IndexFlatL2(dimension)
-    vector_store = FAISS(embedding_function=embeddings, index=index, documents=all_chunks)
-
-    retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 10})
-    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
-
-def upload_book(file, progress=gr.Progress()):
-    with tempfile.NamedTemporaryFile(delete=False) as tmp:
-        tmp.write(file.read())
-        tmp_path = tmp.name
-
-    progress(0.2, desc="تحميل الملف...")
-    extracted_text = extract_text_from_file(tmp_path)
-    uploaded_texts.append(extracted_text)
-    progress(0.5, desc="معالجة النص...")
-
-    train_from_texts(uploaded_texts)
-    progress(1.0, desc="اكتمل التدريب!")
-    return "النظام جاهز للإجابة على أسئلتك"
-
-def answer_question(user_question):
-    if qa_chain is None:
-        return "الرجاء رفع كتاب أولاً."
-    prompt = format_arabic_prompt(user_question)
-    result = qa_chain.run(prompt)
-    return result
+def build_vector_store(chunks):
+    vectors = embedder.encode(chunks)
+    dim = vectors.shape[1]
+    idx = faiss.IndexFlatL2(dim)
+    idx.add(vectors)
+    return idx, vectors
+
+def retrieve_relevant_chunks(question, idx, chunks, vectors):
+    q_vec = embedder.encode([question])
+    D, I = idx.search(q_vec, TOP_K)
+    return [chunks[i] for i in I[0] if i < len(chunks)]
+
+def generate_answer(context_chunks, question):
+    context = " \n".join(context_chunks)
+    prompt = f"سؤال: {question}\nمحتوى ذو صلة: {context}\nجواب:"
+    inputs = gen_tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
+    outputs = gen_model.generate(**inputs, max_new_tokens=100)
+    answer = gen_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    answer = answer.split("جواب:")[-1].strip()
+    return answer
+
+# ====== Gradio Functions ======
+def upload_and_train(files):
+    global index, chunks
+
+    all_text = ""
+    for file in files:
+        suffix = os.path.splitext(file.name)[-1].lower()
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+            tmp.write(file.read())
+            tmp_path = tmp.name
+        if suffix == ".pdf":
+            all_text += extract_text_from_pdf(tmp_path) + "\n"
+        elif suffix in [".docx", ".doc"]:
+            all_text += extract_text_from_docx(tmp_path) + "\n"
+        os.unlink(tmp_path)
+
+    chunks = chunk_text(all_text)
+    index, vectors = build_vector_store(chunks)
+
+    return "✅ النظام جاهز للإجابة على أسئلتك"
+
+def ask_question(user_question):
+    if index is None:
+        return "الرجاء رفع الكتب أولاً وتدريب النظام."
+    rel_chunks = retrieve_relevant_chunks(user_question, index, chunks, None)
+    answer = generate_answer(rel_chunks, user_question)
+    return answer
+
+# ====== Gradio Interface ======
+upload = gr.File(file_types=[".pdf", ".doc", ".docx"], file_count="multiple")
+train_btn = gr.Button("ابدأ التدريب")
+train_output = gr.Textbox()
+question_input = gr.Textbox(placeholder="اكتب سؤالك هنا باللغة العربية")
+answer_output = gr.Textbox()
+ask_btn = gr.Button("أرسل السؤال")
 
 with gr.Blocks() as demo:
-    with gr.Tab("تحميل الكتب"):
-        upload_button = gr.File(label="ارفع كتابك (.pdf .docx .doc)", file_types=[".pdf", ".docx", ".doc"])
-        upload_output = gr.Textbox(label="حالة النظام")
-        upload_button.upload(upload_book, inputs=upload_button, outputs=upload_output)
-
-    with gr.Tab("اسأل الكتاب"):
-        question = gr.Textbox(label="اكتب سؤالك بالعربية")
-        answer = gr.Textbox(label="الإجابة")
-        ask_button = gr.Button("إرسال السؤال")
-        ask_button.click(answer_question, inputs=question, outputs=answer)
-
-demo.launch(share=True)
+    gr.Markdown("# 🧠 محاكاة دماغ المؤلف - نظام ذكي للإجابة على الأسئلة من كتبك بالعربية")
+    upload.render()
+    train_btn.render()
+    train_output.render()
+    question_input.render()
+    ask_btn.render()
+    answer_output.render()
+
+    train_btn.click(upload_and_train, inputs=[upload], outputs=[train_output])
+    ask_btn.click(ask_question, inputs=[question_input], outputs=[answer_output])
+
+# Launch
+if __name__ == "__main__":
+    demo.launch()