Update app.py

app.py

@@ -9,15 +9,16 @@ from huggingface_hub import login
 from transformers import (
     AutoTokenizer, 
     AutoModelForCausalLM, 
-    LlavaNextProcessor, 
-    LlavaNextForConditionalGeneration,
+    # AutoProcessor, # Replaced by ViltProcessor for VQA
+    # AutoModelForVision2Seq, # Replaced by ViltForQuestionAnswering for VQA
     WhisperProcessor,
     WhisperForConditionalGeneration,
-    SpeechT5Processor,
-    SpeechT5ForTextToSpeech,
-    SpeechT5HifiGan
+    ViltProcessor, # Added for ViLT VQA model
+    ViltForQuestionAnswering # Added for ViLT VQA model
 )
-from datasets import load_dataset
+from parler_tts import ParlerTTSForConditionalGeneration
+from transformers import AutoFeatureExtractor
+
 import os
 import scipy.io.wavfile as wavfile
 import io
@@ -29,166 +30,106 @@ token = os.getenv("HF_API_TOKEN")
 if token:
     login(token=token)
 else:
-    print("Warning: HF_API_TOKEN not set. Some models might require login.")
+    print("تحذير: لم يتم تعيين متغير البيئة HF_API_TOKEN. بعض النماذج قد تتطلب تسجيل الدخول.")
 
 # Device Configuration
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
+print(f"استخدام الجهاز: {device}")
 
-# 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}")
+# 1. Text Generation Model: distilgpt2 (Lightweight)
+text_model_name = "distilgpt2"
+print(f"تحميل نموذج النص: {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
+    torch_dtype=torch.float32, 
+    device_map="auto"
 )
-print("Text model loaded.")
-
-# 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(
+if text_tokenizer.pad_token is None:
+    text_tokenizer.pad_token = text_tokenizer.eos_token
+print("تم تحميل نموذج النص.")
+
+# 2. Image Analysis Model: dandelin/vilt-b32-finetuned-vqa (Lightweight, Public VQA)
+image_model_name = "dandelin/vilt-b32-finetuned-vqa"
+print(f"تحميل نموذج الصور (VQA): {image_model_name}")
+image_processor = ViltProcessor.from_pretrained(image_model_name)
+image_model = ViltForQuestionAnswering.from_pretrained(
     image_model_name, 
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32, 
     device_map="auto"
 )
-print("Image model loaded.")
+print("تم تحميل نموذج الصور (VQA).")
 
-# 3. Speech-to-Text Model: openai/whisper-large-v3
-stt_model_name = "openai/whisper-large-v3"
-print(f"Loading STT model: {stt_model_name}")
+# 3. Speech-to-Text Model: openai/whisper-tiny (Lightweight)
+stt_model_name = "openai/whisper-tiny"
+print(f"تحميل نموذج تحويل الكلام إلى نص: {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.")
-
-# 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"
+stt_model.config.forced_decoder_ids = None
+print("تم تحميل نموذج تحويل الكلام إلى نص.")
 
-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)
-
-# 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
-
-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.")
+# 4. Text-to-Speech Model: parler-tts/parler-tts-tiny-v1 (Lightweight)
+tts_model_repo_id = "parler-tts/parler-tts-tiny-v1"
+print(f"تحميل نموذج تحويل النص إلى كلام: {tts_model_repo_id}")
+tts_model = ParlerTTSForConditionalGeneration.from_pretrained(tts_model_repo_id).to(device)
+tts_feature_extractor = AutoFeatureExtractor.from_pretrained(tts_model_repo_id)
+print("تم تحميل مكونات تحويل النص إلى كلام.")
 
 # --- Helper Functions for Model Inference ---
 
-# 1. Text Generation
+# 1. Text Generation (using distilgpt2)
 def generate_text_response(prompt_text):
     try:
-        # 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)
-
+        full_prompt = f"السؤال: {prompt_text}\nالإجابة الودية والواضحة:"
+        inputs = text_tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True, max_length=512).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
+            max_new_tokens=150,
             temperature=0.7,
             top_k=50,
             do_sample=True,
-            pad_token_id=text_tokenizer.eos_token_id
+            pad_token_id=text_tokenizer.eos_token_id,
+            no_repeat_ngram_size=2
         )
-        # 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}")
+        print(f"خطأ في توليد النص: {e}")
         return f"خطأ في معالجة النص: {str(e)}"
 
-# 2. Image Analysis
+# 2. Image Analysis (using dandelin/vilt-b32-finetuned-vqa)
 def analyze_image(pil_image, question_text=None):
     try:
         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
+        else:
             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:"
+            # ViLT is a VQA model, it needs a question.
+            # If no question, we can ask a generic one, or return a message.
+            # For now, let's ask a generic question if none is provided.
+            question_text = "What is in this image?"
+        
+        # Prepare inputs for ViLT
+        encoding = image_processor(pil_image, question_text, return_tensors="pt").to(image_model.device)
+        
+        # Forward pass
+        with torch.no_grad():
+            outputs = image_model(**encoding)
+        
+        logits = outputs.logits
+        idx = logits.argmax(-1).item()
+        response_text = image_model.config.id2label[idx]
         
-        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}")
+        print(f"خطأ في تحليل الصورة: {e}")
         return f"خطأ في تحليل الصورة: {str(e)}"
 
-# 3. Audio Processing (STT and TTS)
+# 3. Audio Processing (STT with Whisper Tiny and TTS with ParlerTTS Tiny)
 def process_audio(audio_input):
     try:
         if audio_input is None:
@@ -196,41 +137,34 @@ def process_audio(audio_input):
             
         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
+        else:
             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()
+        predicted_ids = stt_model.generate(input_features, language="ar")
+        transcription = stt_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0].strip()
 
         if not transcription:
             return "لم يتمكن النموذج من استخراج نص من الصوت.", "", (16000, np.array([], dtype=np.int16))
 
-        # Text Generation (Response to transcription)
         text_response = generate_text_response(transcription)
-
-        # 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 
+        prompt = text_response
+        with torch.no_grad():
+            generation_output = tts_model.generate(input_ids=None, 
+                                                   prompt=prompt, 
+                                                   do_sample=True, 
+                                                   temperature=1.0).cpu().numpy().squeeze()
+        audio_output_np = generation_output
+        tts_sample_rate = tts_model.config.sampling_rate
 
         return transcription, text_response, (tts_sample_rate, audio_output_np)
     except Exception as e:
-        print(f"Error in audio processing: {e}")
-        # Return empty audio with a valid sample rate for Gradio
+        print(f"خطأ في معالجة الصوت: {e}")
         empty_audio_data = np.array([], dtype=np.float32)
         return f"خطأ في معالجة الصوت: {str(e)}", "", (16000, empty_audio_data)
 
@@ -239,7 +173,7 @@ def process_file(file_obj):
     try:
         if file_obj is None:
             return "الرجاء رفع ملف أولاً."
-        file_path = file_obj.name # Gradio File object has a .name attribute with the temp path
+        file_path = file_obj.name
         text_content = ""
         if file_path.endswith(".pdf"):
             with fitz.open(file_path) as doc:
@@ -256,23 +190,21 @@ def process_file(file_obj):
         if not text_content.strip():
             return "الملف فارغ أو لا يمكن قراءة محتواه النصي."
 
-        # 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
+        max_context_len = 1000
         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}")
+        print(f"خطأ في معالجة الملف: {e}")
         return f"خطأ في قراءة الملف: {str(e)}"
 
 # --- 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🎙️ تحويل الصوت إلى نص والرد صوتياً")
+    gr.Markdown("# 🤖 Kemo Chat V3.2 - مساعد ذكي متعدد الوسائط (نماذج خفيفة الوزن - ViLT VQA)")
+    gr.Markdown("🎯 تفاعل معي عبر النصوص، الصور، الصوت أو الملفات! (باستخدام نماذج أقل استهلاكًا للذاكرة).")
+    gr.Markdown("📁 يدعم الملفات: PDF، Excel، CSV\n🖼️ يدعم الإجابة على الأسئلة حول الصور (VQA)\n🎙️ تحويل الصوت إلى نص والرد صوتياً")
 
     with gr.Tab("💬 المحادثة النصية"):
         text_input = gr.Textbox(label="اكتب سؤالك أو رسالتك هنا", lines=3)
@@ -280,23 +212,23 @@ with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: A
         text_submit = gr.Button("إرسال", variant="primary")
         text_submit.click(fn=generate_text_response, inputs=text_input, outputs=text_output)
 
-    with gr.Tab("🖼️ تحليل الصور"):
-        gr.Markdown("ارفع صورة واسأل عنها أو اطلب وصفًا لها.")
+    with gr.Tab("🖼️ تحليل الصور (سؤال وجواب)"):
+        gr.Markdown("ارفع صورة واطرح سؤالاً عنها.")
         with gr.Row():
-            image_input = gr.Image(type="pil", label="ارفع صورة") # type="pil" for PIL Image object
+            image_input = gr.Image(type="pil", label="ارفع صورة")
             with gr.Column():
-                image_question = gr.Textbox(label="اختياري: اسأل سؤال عن الصورة", lines=2)
-                image_output = gr.Textbox(label="الوصف أو الإجابة", lines=5, interactive=False)
+                image_question = gr.Textbox(label="اطرح سؤالاً عن الصورة (مطلوب لـ ViLT)", lines=2, placeholder="مثال: What color is the car?")
+                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("🎤 التفاعل الصوتي"):
         gr.Markdown("سجّل رسالة صوتية، سيتم تحويلها إلى نص، ثم الرد عليها نصيًا وصوتيًا.")
-        audio_input = gr.Audio(sources=["microphone"], type="numpy", label="سجّل رسالتك الصوتية") # type="numpy" for (sr, data)
+        audio_input = gr.Audio(sources=["microphone"], type="numpy", label="سجّل رسالتك الصوتية")
         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_output_player = gr.Audio(label="الرد الصوتي من المساعد", type="numpy", interactive=False)
         audio_submit = gr.Button("معالجة الصوت", variant="primary")
         audio_submit.click(fn=process_audio, 
                            inputs=audio_input, 
@@ -310,6 +242,6 @@ with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: A
         file_submit.click(fn=process_file, inputs=file_input, outputs=file_output)
 
 if __name__ == "__main__":
-    print("Launching Gradio Demo...")
-    demo.launch(share=True) # share=True for public link, remove if not needed
+    print("Launching Gradio Demo (Lightweight Models with ViLT VQA)...")
+    demo.launch(share=True)