Update app.py

app.py

@@ -1,136 +1,107 @@
+import os
+import re
+import json
 import gradio as gr
+from PIL import Image
+
 import torch
-import os
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import spaces
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
+import spaces  # 👈 Hugging Face ZeroGPU
 
-# 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.")
+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"))
 
-# Check for GPU support and configure appropriately
+# ---- Device selection ----
 device = "cuda" if torch.cuda.is_available() else "cpu"
-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(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
+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.")
+    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])
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True) 
+    demo.launch()