Update app.py

app.py

@@ -1,107 +1,136 @@
-import os
-import re
-import json
 import gradio as gr
-from PIL import Image
-
 import torch
-from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
-import spaces  # 👈 Hugging Face ZeroGPU
+import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import spaces
 
-MODEL_NAME = os.environ.get("MODEL_NAME", "NAMAA-Space/Qari-OCR-0.1-VL-2B-Instruct")
-MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "1024"))
+# Load Hugging Face token from the environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")
+if HF_TOKEN is None:
+    raise ValueError("HF_TOKEN environment variable is not set. Please set it before running the script.")
 
-# ---- Device selection ----
+# Check for GPU support and configure appropriately
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if device == "cuda" else torch.float32
-
-print(f"Device being used: {device}")
-
-# ---- Load model & processor ----
-processor = AutoProcessor.from_pretrained(MODEL_NAME)
-
-model = Qwen2VLForConditionalGeneration.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=dtype,
-    device_map="auto" if device == "cuda" else {"": "cpu"},
-)
-
-print("Model loaded successfully!")
-
-def _mk_messages(image_path: str, prompt_info: str):
-    return [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": f"file://{image_path}"},
-                {"type": "text", "text": f"""
-You are an advanced invoice OCR system...
-(extractions instructions same as notebook)
-Extra hints from user: {prompt_info}
-""".strip()},
-            ],
-        }
-    ]
-
-def _extract_json(text: str):
-    text = text.strip()
-    if text.startswith("{") and text.endswith("}"):
-        try:
-            return json.loads(text)
-        except Exception:
-            pass
-    m = re.search(r"\{[\s\S]*\}", text)
-    if m:
-        block = m.group(0)
-        try:
-            return json.loads(block)
-        except Exception:
-            pass
-    return {"other_text": text}
-
-@spaces.GPU(duration=120)  # 👈 Request ZeroGPU for 2 minutes
-def infer(image: Image.Image, prompt_info: str):
-    if image is None:
-        return "Please upload an image.", {}
-
-    tmp_path = "input_image.png"
-    image.save(tmp_path)
-
-    messages = _mk_messages(tmp_path, prompt_info)
-    chat_text = processor.apply_chat_template(messages, add_generation_prompt=True)
-
-    inputs = processor(
-        text=[chat_text],
-        images=[Image.open(tmp_path)],
-        return_tensors="pt",
-    )
-    inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
-
-    with torch.no_grad():
-        generated_ids = model.generate(
-            **inputs,
-            max_new_tokens=MAX_NEW_TOKENS,
-            do_sample=False,
-        )
-
-    gen_only = generated_ids[:, inputs["input_ids"].shape[1]:]
-    text_out = processor.batch_decode(gen_only, skip_special_tokens=True)[0].strip()
-    parsed = _extract_json(text_out)
-
-    return text_out, parsed
-
-with gr.Blocks(title="Qari OCR (ZeroGPU)") as demo:
-    gr.Markdown("# Qari OCR · ZeroGPU\nUpload an invoice image and (optionally) add extraction hints.")
+zero = torch.Tensor([0]).to(device)
+print(f"Device being used: {zero.device}")
+
+# Model configurations
+MSA_TO_SYRIAN_MODEL = "Omartificial-Intelligence-Space/Shami-MT"
+SYRIAN_TO_MSA_MODEL = "Omartificial-Intelligence-Space/SHAMI-MT-2MSA"
+
+# Load models and tokenizers
+print("Loading MSA to Syrian model...")
+msa_to_syrian_tokenizer = AutoTokenizer.from_pretrained(MSA_TO_SYRIAN_MODEL)
+msa_to_syrian_model = AutoModelForSeq2SeqLM.from_pretrained(MSA_TO_SYRIAN_MODEL).to(device)
+
+print("Loading Syrian to MSA model...")
+syrian_to_msa_tokenizer = AutoTokenizer.from_pretrained(SYRIAN_TO_MSA_MODEL)
+syrian_to_msa_model = AutoModelForSeq2SeqLM.from_pretrained(SYRIAN_TO_MSA_MODEL).to(device)
+
+print("Models loaded successfully!")
+
+@spaces.GPU(duration=120)
+def translate_msa_to_syrian(text):
+    """Translate from Modern Standard Arabic to Syrian dialect"""
+    if not text.strip():
+        return ""
+    
+    try:
+        input_ids = msa_to_syrian_tokenizer(text, return_tensors="pt").input_ids.to(device)
+        outputs = msa_to_syrian_model.generate(input_ids, max_length=128, num_beams=5, early_stopping=True)
+        translated_text = msa_to_syrian_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return translated_text
+    except Exception as e:
+        return f"Translation error: {str(e)}"
+
+@spaces.GPU(duration=120)
+def translate_syrian_to_msa(text):
+    """Translate from Syrian dialect to Modern Standard Arabic"""
+    if not text.strip():
+        return ""
+    
+    try:
+        input_ids = syrian_to_msa_tokenizer(text, return_tensors="pt").input_ids.to(device)
+        outputs = syrian_to_msa_model.generate(input_ids, max_length=128, num_beams=5, early_stopping=True)
+        translated_text = syrian_to_msa_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return translated_text
+    except Exception as e:
+        return f"Translation error: {str(e)}"
+
+def bidirectional_translate(text, direction):
+    """Handle bidirectional translation based on user selection"""
+    if direction == "MSA → Syrian":
+        return translate_msa_to_syrian(text)
+    elif direction == "Syrian → MSA":
+        return translate_syrian_to_msa(text)
+    else:
+        return "Please select a translation direction"
+
+# Create Gradio interface
+with gr.Blocks(title="SHAMI-MT: Bidirectional Syria Arabic Dialect MT Framework") as demo:
+    
+    gr.HTML("""
+    <div style="text-align: center; margin-bottom: 2rem;">
+        <h1>🌍 SHAMI-MT: Bidirectional Arabic Translation</h1>
+        <p>Translate between Modern Standard Arabic (MSA) and Syrian Dialect</p>
+        <p><strong>Built on AraT5v2-base-1024 architecture</strong></p>
+    </div>
+    """)
+    
     with gr.Row():
-        with gr.Column():
-            img_in = gr.Image(type="pil", label="Invoice Image")
-            prompt_box = gr.Textbox(label="Extra hints (optional)")
-            run_btn = gr.Button("Run OCR")
-        with gr.Column():
-            txt_out = gr.Textbox(label="Raw Model Output", lines=10)
-            json_out = gr.JSON(label="Parsed JSON")
-
-    run_btn.click(infer, inputs=[img_in, prompt_box], outputs=[txt_out, json_out])
+        with gr.Column(scale=1):
+            gr.HTML("""
+            <div style="background: #f8f9fa; padding: 1rem; border-radius: 8px; margin: 1rem 0;">
+                <h3>📚 Model Information</h3>
+                <ul>
+                    <li><strong>Model Type:</strong> Sequence-to-Sequence Translation</li>
+                    <li><strong>Base Model:</strong> UBC-NLP/AraT5v2-base-1024</li>
+                    <li><strong>Languages:</strong> Arabic (MSA ↔ Syrian Dialect)</li>
+                    <li><strong>Device:</strong> GPU/CPU Auto-detection</li>
+                </ul>
+            </div>
+            """)
+    
+        with gr.Column(scale=2):
+            direction = gr.Dropdown(
+                choices=["MSA → Syrian", "Syrian → MSA"],
+                value="MSA → Syrian",
+                label="Translation Direction"
+            )
+            
+            input_text = gr.Textbox(
+                label="Input Text",
+                placeholder="Enter Arabic text here...",
+                lines=5
+            )
+            
+            translate_btn = gr.Button("🚀 Translate", variant="primary")
+            
+            output_text = gr.Textbox(
+                label="Translation",
+                lines=5
+            )
+    
+    # Connect the interface
+    translate_btn.click(
+        fn=bidirectional_translate,
+        inputs=[input_text, direction],
+        outputs=output_text
+    )
+    
+    # Add example inputs
+    gr.Examples(
+        examples=[
+            ["أنا لا أعرف إذا كان سيتمكن من الحضور اليوم أم لا.", "MSA → Syrian"],
+            ["كيف حالك؟", "MSA → Syrian"],
+            ["ما بعرف إذا رح يقدر يجي اليوم ولا لأ.", "Syrian → MSA"],
+            ["شلونك؟", "Syrian → MSA"]
+        ],
+        inputs=[input_text, direction],
+        outputs=output_text,
+        fn=bidirectional_translate
+    )
 
+# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)