Update app.py

app.py

@@ -1,136 +1,315 @@
+# coding: utf-8
 import gradio as gr
 import torch
-from TTS.api import TTS
 import numpy as np
 from PIL import Image
 import fitz  # PyMuPDF
 import pandas as pd
 from huggingface_hub import login
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    LlavaNextProcessor, 
+    LlavaNextForConditionalGeneration,
+    WhisperProcessor,
+    WhisperForConditionalGeneration,
+    SpeechT5Processor,
+    SpeechT5ForTextToSpeech,
+    SpeechT5HifiGan
+)
+from datasets import load_dataset
 import os
+import scipy.io.wavfile as wavfile
+import io
 
-# تسجيل الدخول باستخدام المفتاح من المتغير البيئي
+# --- Configuration & Model Loading ---
+
+# Hugging Face Hub Login
 token = os.getenv("HF_API_TOKEN")
 if token:
     login(token=token)
 else:
-    raise ValueError("لم يتم تعيين المفتاح البيئي HF_API_TOKEN.")
+    print("Warning: HF_API_TOKEN not set. Some models might require login.")
+
+# Device Configuration
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
-# اسم النموذج
-mistral_model_name = "TheBloke/Mistral-7B-Instruct-v0.1-GGUF"
+# 1. Text Generation Model: microsoft/Phi-4-reasoning-plus
+text_model_name = "microsoft/Phi-4-reasoning-plus"
+print(f"Loading text model: {text_model_name}")
+text_tokenizer = AutoTokenizer.from_pretrained(text_model_name)
+text_model = AutoModelForCausalLM.from_pretrained(
+    text_model_name, 
+    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32, # bfloat16 for better performance if available
+    device_map="auto", # Automatically uses available GPUs or CPU
+    trust_remote_code=True # Phi models may require this
+)
+print("Text model loaded.")
 
-# تحميل التوكنيزر والموديل
-mistral_tokenizer = AutoTokenizer.from_pretrained(mistral_model_name)
-mistral_model = AutoModelForCausalLM.from_pretrained(mistral_model_name)
+# 2. Image Analysis Model: llava-hf/llava-1.5-7b-hf
+image_model_name = "llava-hf/llava-1.5-7b-hf"
+print(f"Loading image model: {image_model_name}")
+image_processor = LlavaNextProcessor.from_pretrained(image_model_name)
+image_model = LlavaNextForConditionalGeneration.from_pretrained(
+    image_model_name, 
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto"
+)
+print("Image model loaded.")
 
-# Initialize device
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# 3. Speech-to-Text Model: openai/whisper-large-v3
+stt_model_name = "openai/whisper-large-v3"
+print(f"Loading STT model: {stt_model_name}")
+stt_processor = WhisperProcessor.from_pretrained(stt_model_name)
+stt_model = WhisperForConditionalGeneration.from_pretrained(stt_model_name).to(device)
+stt_model.config.forced_decoder_ids = None # For multilingual capabilities
+print("STT model loaded.")
 
-# Text model: Mistral 7B
-mistral_model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-mistral_tokenizer = AutoTokenizer.from_pretrained(mistral_model_name)
-mistral_model = AutoModelForCausalLM.from_pretrained(mistral_model_name, torch_dtype=torch.float16 if device == "cuda" else torch.float32).to(device)
+# 4. Text-to-Speech Model: microsoft/speecht5_tts & microsoft/speecht5_hifigan
+tts_processor_name = "microsoft/speecht5_tts"
+tts_model_name = "microsoft/speecht5_tts"
+tts_vocoder_name = "microsoft/speecht5_hifigan"
 
-# Image model: BLIP-2
-blip_model_name = "Salesforce/blip2-opt-2.7b"
-blip_processor = Blip2Processor.from_pretrained(blip_model_name)
-blip_model = Blip2ForConditionalGeneration.from_pretrained(blip_model_name, torch_dtype=torch.float16 if device == "cuda" else torch.float32).to(device)
+print(f"Loading TTS processor: {tts_processor_name}")
+tts_processor = SpeechT5Processor.from_pretrained(tts_processor_name)
+print(f"Loading TTS model: {tts_model_name}")
+tts_model = SpeechT5ForTextToSpeech.from_pretrained(tts_model_name).to(device)
+print(f"Loading TTS vocoder: {tts_vocoder_name}")
+tts_vocoder = SpeechT5HifiGan.from_pretrained(tts_vocoder_name).to(device)
 
-# Speech-to-text: Whisper
-whisper_model_name = "openai/whisper-small"
-whisper_processor = WhisperProcessor.from_pretrained(whisper_model_name)
-whisper_model = WhisperForConditionalGeneration.from_pretrained(whisper_model_name).to(device)
+# Load speaker embeddings (example from CMU ARCTIC dataset)
+print("Loading speaker embeddings for TTS...")
+try:
+    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+    # speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0).to(device) # A specific speaker
+    # Let's pick a female voice, e.g., slt (female)
+    # Find an entry for 'slt' speaker, or use a default if not easily found by index.
+    # For simplicity, using a known good index for a female voice if available, otherwise a default.
+    # This part might need adjustment based on the dataset structure or desired voice.
+    # Example: embeddings_dataset is a list of dicts. Find one with speaker 'slt'.
+    # For now, let's use a fixed index that often corresponds to a clear female voice.
+    speaker_id_index = 7306 # Example index, might need to find a better one or allow selection
+    # A common female speaker from this dataset is 'slt'. Let's try to find her or use a default.
+    # The dataset has 'slt_arctic_a0001' to 'slt_arctic_b0593'. Index 7306 is 'awb_arctic_a0001'.
+    # Let's try to find 'slt' if possible, otherwise use a default. The dataset has 8803 entries.
+    # For simplicity, we'll use a default index. A common one used in examples is 7306.
+    # A known female voice index from this dataset (often 'slt' or similar) is around this range.
+    # Let's use a specific one for 'slt' if we can find it, otherwise a default.
+    # The dataset items are like {'audio': ..., 'filename': 'arctic_a0001.wav', 'speaker': 'slt', 'text': ..., 'xvector': ...}
+    # Let's try to find an 'slt' speaker explicitly
+    slt_speaker_embedding = None
+    for i in range(len(embeddings_dataset)):
+        if embeddings_dataset[i]['speaker'] == 'slt':
+            slt_speaker_embedding = torch.tensor(embeddings_dataset[i]["xvector"]).unsqueeze(0).to(device)
+            print("Found 'slt' speaker embedding.")
+            break
+    if slt_speaker_embedding is None:
+        print("Could not find 'slt' speaker, using default index 7306.")
+        slt_speaker_embedding = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0).to(device)
+    speaker_embeddings = slt_speaker_embedding
 
-# Text-to-speech: Coqui TTS
-tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False)
+except Exception as e:
+    print(f"Could not load speaker embeddings: {e}. TTS might not work optimally or at all.")
+    speaker_embeddings = torch.randn((1, 512)).to(device) # Fallback to random embeddings
+print("TTS components loaded.")
 
-# === Text generation ===
-def generate_text_response(prompt):
+# --- Helper Functions for Model Inference ---
+
+# 1. Text Generation
+def generate_text_response(prompt_text):
     try:
-        friendly_prompt = f"أجب على السؤال التالي بطريقة ودية وواضحة:\n{prompt}"
-        inputs = mistral_tokenizer(friendly_prompt, return_tensors="pt").to(device)
-        outputs = mistral_model.generate(**inputs, max_length=200, num_return_sequences=1)
-        return mistral_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Phi-4-reasoning-plus uses a specific chat template
+        # System prompt from the model card (simplified for general use here)
+        system_prompt = "You are a helpful AI assistant. Please provide a clear and friendly answer."
+        # The model card suggests a more complex system prompt for reasoning tasks.
+        # For a general chatbot, a simpler one might be fine, or adjust as needed.
+        # Original prompt: "أجب على السؤال التالي بطريقة ودية وواضحة:"
+        # We will use the user's prompt directly with a general system prompt.
+        
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"أجب على السؤال التالي بطريقة ودية وواضحة: {prompt_text}"}
+        ]
+        
+        # Tokenize the chat
+        # The model card for Phi-4-reasoning-plus shows: inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+        # However, apply_chat_template might not be available or work the same for all AutoTokenizers directly without specific setup.
+        # A more general approach for models expecting chat format:
+        prompt_for_model = text_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = text_tokenizer(prompt_for_model, return_tensors="pt", padding=True, truncation=True, max_length=2048).to(text_model.device)
+
+        outputs = text_model.generate(
+            **inputs,
+            max_new_tokens=300, # Increased from 200 for potentially longer answers
+            # Parameters from Phi-4 model card for reasoning, can be adjusted for chat
+            # temperature=0.8, 
+            # top_p=0.95, 
+            # do_sample=True,
+            # For more deterministic chat, we might use different settings
+            temperature=0.7,
+            top_k=50,
+            do_sample=True,
+            pad_token_id=text_tokenizer.eos_token_id
+        )
+        # Decode, skipping the prompt part
+        response_text = text_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        return response_text.strip()
     except Exception as e:
+        print(f"Error in text generation: {e}")
         return f"خطأ في معالجة النص: {str(e)}"
 
-# === Image analysis ===
-def analyze_image(image, question=None):
+# 2. Image Analysis
+def analyze_image(pil_image, question_text=None):
     try:
-        image = Image.fromarray(image).convert("RGB")
-        prompt = question if question else "صف محتوى الصورة بالتفصيل"
-        inputs = blip_processor(images=image, text=prompt, return_tensors="pt").to(device)
-        outputs = blip_model.generate(**inputs)
-        return blip_processor.decode(outputs[0], skip_special_tokens=True)
+        if pil_image is None:
+            return "الرجاء رفع صورة أولاً."
+        # Convert numpy array from Gradio to PIL Image if necessary
+        if isinstance(pil_image, np.ndarray):
+            pil_image = Image.fromarray(pil_image).convert("RGB")
+        else: # Assuming it's already a PIL image path or object from gr.Image
+            pil_image = pil_image.convert("RGB")
+
+        if not question_text or question_text.strip() == "":
+            prompt_for_model = "USER: <image>\nصف محتوى الصورة بالتفصيل\nASSISTANT:"
+        else:
+            prompt_for_model = f"USER: <image>\n{question_text}\nASSISTANT:"
+        
+        inputs = image_processor(text=prompt_for_model, images=pil_image, return_tensors="pt").to(image_model.device)
+        outputs = image_model.generate(**inputs, max_new_tokens=100)
+        response_text = image_processor.decode(outputs[0], skip_special_tokens=True)
+        # LLaVA response often includes the prompt, so we might need to clean it.
+        # Typically, the response starts after "ASSISTANT: "
+        assistant_marker = "ASSISTANT:"
+        if assistant_marker in response_text:
+            response_text = response_text.split(assistant_marker, 1)[-1].strip()
+        return response_text
     except Exception as e:
+        print(f"Error in image analysis: {e}")
         return f"خطأ في تحليل الصورة: {str(e)}"
 
-# === Audio processing ===
-def process_audio(audio):
+# 3. Audio Processing (STT and TTS)
+def process_audio(audio_input):
     try:
-        sample_rate, audio_data = audio
-        audio_data = audio_data.astype(np.float32) / np.max(np.abs(audio_data))
-        input_features = whisper_processor(audio_data, sampling_rate=sample_rate, return_tensors="pt").input_features.to(device)
-        predicted_ids = whisper_model.generate(input_features)
-        transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+        if audio_input is None:
+            return "الرجاء تسجيل الصوت أولاً.", "", (16000, np.array([], dtype=np.int16))
+            
+        sample_rate, audio_data = audio_input
+        
+        # Ensure audio_data is float32 and normalized for Whisper
+        if audio_data.dtype != np.float32:
+            audio_data = audio_data.astype(np.float32)
+        if np.max(np.abs(audio_data)) > 0:
+            audio_data = audio_data / np.max(np.abs(audio_data))
+        else: # Handle silent audio
+            return "تم استقبال صوت صامت.", "", (16000, np.array([], dtype=np.int16))
+
+        # Speech-to-Text
+        input_features = stt_processor(audio_data, sampling_rate=sample_rate, return_tensors="pt").input_features.to(device)
+        # Generate token ids
+        predicted_ids = stt_model.generate(input_features, language="ar") # Specify language for better accuracy if known
+        # Decode token ids to text
+        transcription = stt_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+        transcription = transcription.strip()
+
+        if not transcription:
+            return "لم يتمكن النموذج من استخراج نص من الصوت.", "", (16000, np.array([], dtype=np.int16))
+
+        # Text Generation (Response to transcription)
         text_response = generate_text_response(transcription)
-        audio_output = tts.tts(text_response)
-        audio_output = np.array(audio_output, dtype=np.float32)
-        return transcription, text_response, (16000, audio_output)
+
+        # Text-to-Speech
+        inputs = tts_processor(text=text_response, return_tensors="pt")
+        speech_values = tts_model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=tts_vocoder)
+        
+        # Ensure output is a NumPy array and handle sample rate
+        audio_output_np = speech_values.cpu().numpy()
+        # SpeechT5 typically outputs at 16kHz
+        tts_sample_rate = 16000 
+
+        return transcription, text_response, (tts_sample_rate, audio_output_np)
     except Exception as e:
-        return f"خطأ في معالجة الصوت: {str(e)}", "", None
+        print(f"Error in audio processing: {e}")
+        # Return empty audio with a valid sample rate for Gradio
+        empty_audio_data = np.array([], dtype=np.float32)
+        return f"خطأ في معالجة الصوت: {str(e)}", "", (16000, empty_audio_data)
 
-# === File processing ===
-def process_file(file):
+# 4. File Processing
+def process_file(file_obj):
     try:
-        if file.name.endswith(".pdf"):
-            with fitz.open(file.name) as doc:
-                text = "\n".join(page.get_text() for page in doc)
-        elif file.name.endswith((".xlsx", ".xls")):
-            df = pd.read_excel(file.name)
-            text = df.to_string()
-        elif file.name.endswith(".csv"):
-            df = pd.read_csv(file.name)
-            text = df.to_string()
+        if file_obj is None:
+            return "الرجاء رفع ملف أولاً."
+        file_path = file_obj.name # Gradio File object has a .name attribute with the temp path
+        text_content = ""
+        if file_path.endswith(".pdf"):
+            with fitz.open(file_path) as doc:
+                text_content = "\n".join(page.get_text() for page in doc)
+        elif file_path.endswith((".xlsx", ".xls")):
+            df = pd.read_excel(file_path)
+            text_content = df.to_string()
+        elif file_path.endswith(".csv"):
+            df = pd.read_csv(file_path)
+            text_content = df.to_string()
         else:
-            return "❌ نوع الملف غير مدعوم حالياً."
+            return "❌ نوع الملف غير مدعوم حالياً (يدعم PDF, Excel, CSV)."
+
+        if not text_content.strip():
+            return "الملف فارغ أو لا يمكن قراءة محتواه النصي."
 
-        response = generate_text_response(f"الملف يحتوي على:\n{text}\n\nلخص المحتوى.")
+        # Summarize or query the content
+        # Limit context size for the model if necessary
+        max_context_len = 1500 # Adjust based on model limits and typical file sizes
+        if len(text_content) > max_context_len:
+            text_content = text_content[:max_context_len] + "... [المحتوى تم اختصاره]"
+            
+        response = generate_text_response(f"لخص المحتوى التالي من الملف: \n\n{text_content}")
         return response
     except Exception as e:
+        print(f"Error processing file: {e}")
         return f"خطأ في قراءة الملف: {str(e)}"
 
-# === Gradio Interface ===
-with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: Arial; color: #333; padding: 20px;}") as demo:
-    gr.Markdown("# 🤖 Kemo Chat - مساعد ذكي متعدد الوسائط")
-    gr.Markdown("🎯 تفاعل معي عبر النصوص، الصور، الصوت أو الملفات! يدعم العربية والإنجليزية.")
-    gr.Markdown("📁 يدعم الملفات: PDF، Excel، CSV\n🖼️ يدعم الوصف الذكي للصور\n🎙️ تحويل الصوت إلى نص والرد صوتياً")
+# --- Gradio Interface ---
+with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: Arial; color: #333; padding: 20px;}", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🤖 Kemo Chat V2 - مساعد ذكي متعدد الوسائط (بأدوات Hugging Face الأحدث)")
+    gr.Markdown("🎯 تفاعل معي عبر النصوص، الصور، الصوت أو الملفات! يدعم العربية والإنجليزية (النماذج قد تتفاوت في دعم اللغات).")
+    gr.Markdown("📁 يدعم الملفات: PDF، Excel، CSV\n🖼️ يدعم الوصف الذكي للصور والسؤال عنها\n🎙️ تحويل الصوت إلى نص والرد صوتياً")
 
     with gr.Tab("💬 المحادثة النصية"):
-        text_input = gr.Textbox(label="اكتب سؤالك أو رسالتك")
-        text_output = gr.Textbox(label="الرد")
-        text_submit = gr.Button("إرسال")
+        text_input = gr.Textbox(label="اكتب سؤالك أو رسالتك هنا", lines=3)
+        text_output = gr.Textbox(label="الرد", lines=5, interactive=False)
+        text_submit = gr.Button("إرسال", variant="primary")
         text_submit.click(fn=generate_text_response, inputs=text_input, outputs=text_output)
 
     with gr.Tab("🖼️ تحليل الصور"):
-        image_input = gr.Image(label="ارفع صورة")
-        image_question = gr.Textbox(label="اختياري: اسأل سؤال عن الصورة")
-        image_output = gr.Textbox(label="الوصف أو الإجابة")
-        image_submit = gr.Button("تحليل الصورة")
+        gr.Markdown("ارفع صورة واسأل عنها أو اطلب وصفًا لها.")
+        with gr.Row():
+            image_input = gr.Image(type="pil", label="ارفع صورة") # type="pil" for PIL Image object
+            with gr.Column():
+                image_question = gr.Textbox(label="اختياري: اسأل سؤال عن الصورة", lines=2)
+                image_output = gr.Textbox(label="الوصف أو الإجابة", lines=5, interactive=False)
+        image_submit = gr.Button("تحليل الصورة", variant="primary")
         image_submit.click(fn=analyze_image, inputs=[image_input, image_question], outputs=image_output)
 
     with gr.Tab("🎤 التفاعل الصوتي"):
-        audio_input = gr.Audio(source="microphone", label="سجّل رسالتك")
-        audio_transcription = gr.Textbox(label="النص المستخرج")
-        audio_text_response = gr.Textbox(label="الرد النصي")
-        audio_output = gr.Audio(label="الرد الص��تي")
-        audio_submit = gr.Button("معالجة الصوت")
-        audio_submit.click(fn=process_audio, inputs=audio_input, outputs=[audio_transcription, audio_text_response, audio_output])
+        gr.Markdown("سجّل رسالة صوتية، سيتم تحويلها إلى نص، ثم الرد عليها نصيًا وصوتيًا.")
+        audio_input = gr.Audio(sources=["microphone"], type="numpy", label="سجّل رسالتك الصوتية") # type="numpy" for (sr, data)
+        with gr.Row():
+            audio_transcription = gr.Textbox(label="النص المستخرج من صوتك", interactive=False, lines=2)
+            audio_text_response = gr.Textbox(label="الرد النصي على رسالتك", interactive=False, lines=3)
+        audio_output_player = gr.Audio(label="الرد الصوتي من المساعد", type="numpy", interactive=False) # type="numpy" for (sr, data)
+        audio_submit = gr.Button("معالجة الصوت", variant="primary")
+        audio_submit.click(fn=process_audio, 
+                           inputs=audio_input, 
+                           outputs=[audio_transcription, audio_text_response, audio_output_player])
 
     with gr.Tab("📄 تحليل الملفات"):
-        file_input = gr.File(label="ارفع ملفك (PDF, Excel, CSV)")
-        file_output = gr.Textbox(label="الرد على محتوى الملف")
-        file_submit = gr.Button("تحليل الملف")
+        gr.Markdown("ارفع ملف (PDF, Excel, CSV) وسأقوم بتلخيص محتواه أو الإجابة على أسئلتك حوله.")
+        file_input = gr.File(label="ارفع ملفك (PDF, Excel, CSV)", file_types=[".pdf", ".xls", ".xlsx", ".csv"])
+        file_output = gr.Textbox(label="الرد على محتوى الملف", lines=5, interactive=False)
+        file_submit = gr.Button("تحليل الملف", variant="primary")
         file_submit.click(fn=process_file, inputs=file_input, outputs=file_output)
 
-demo.launch()
+if __name__ == "__main__":
+    print("Launching Gradio Demo...")
+    demo.launch(share=True) # share=True for public link, remove if not needed
+