.

README.md

@@ -8,6 +8,7 @@ sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 license: mit
+short_description: Multilingual document layout parsing with OCR
 models:
   - rednote-hilab/DotsOCR
 tags:

app.py

@@ -27,13 +27,23 @@ def load_model():
     global model, processor
     if model is None:
         print(f"Loading model weights from {MODEL_PATH}...")
+        
+        # Try to use FlashAttention2 if available, otherwise use default attention
+        try:
+            import flash_attn
+            attn_implementation = "flash_attention_2"
+            print("Using FlashAttention2 for faster inference")
+        except ImportError:
+            attn_implementation = "eager"
+            print("FlashAttention2 not available, using default attention")
+        
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_PATH,
-            attn_implementation="flash_attention_2",
-            torch_dtype=torch.bfloat16,
+            dtype=torch.bfloat16,
             device_map="auto",
             trust_remote_code=True,
-            token=HF_TOKEN
+            token=HF_TOKEN,
+            attn_implementation=attn_implementation
         )
         print("Model loaded successfully.")