Download checkpoint from HF hub in OcrReorderPipeline

inference.py

@@ -3,12 +3,22 @@ from transformers import Pipeline
 from PIL import Image
 import base64
 from io import BytesIO
+from huggingface_hub import hf_hub_download
+
+# point at your HF model repo
+HF_MODEL_REPO = "Uddipan107/ocr-layoutlmv3-base-t5-small"
 
 class OcrReorderPipeline(Pipeline):
     def __init__(self, model, tokenizer, processor, device=0):
         super().__init__(model=model, tokenizer=tokenizer,
                          feature_extractor=processor, device=device)
-        proj_state = torch.load("pytorch_model.bin", map_location="cpu")["projection"]
+
+        # ── Download your fine-tuned checkpoint ───────────────────────────
+        ckpt_path = hf_hub_download(repo_id=HF_MODEL_REPO, filename="pytorch_model.bin")
+        ckpt      = torch.load(ckpt_path, map_location="cpu")
+        proj_state= ckpt["projection"]
+
+        # ── Rebuild & load your projection head ────────────────────────────
         self.projection = torch.nn.Sequential(
             torch.nn.Linear(768, model.config.d_model),
             torch.nn.LayerNorm(model.config.d_model),
@@ -31,17 +41,20 @@ class OcrReorderPipeline(Pipeline):
     def _forward(self, model_inputs):
         pv, ids, mask, bbox = (
             model_inputs[k].to(self.device)
-            for k in ("pixel_values","input_ids","attention_mask","bbox")
+            for k in ("pixel_values", "input_ids", "attention_mask", "bbox")
         )
+
         vision_out = self.model.vision_model(
             pixel_values=pv,
             input_ids=ids,
             attention_mask=mask,
             bbox=bbox
         )
+
         seq_len    = ids.size(1)
         text_feats = vision_out.last_hidden_state[:, :seq_len, :]
         proj_feats = self.projection(text_feats)
+
         gen_ids = self.model.text_model.generate(
             inputs_embeds=proj_feats,
             attention_mask=mask,