Update README.md

README.md

@@ -229,10 +229,8 @@ ColQwen3 generalizes to short videos while learning from image-text retrieval ta
 
 ```python
 from pathlib import Path
-from typing import Any, Dict
 
 import torch
-from torch.nn.utils.rnn import pad_sequence
 from transformers import AutoModel, AutoProcessor
 
 MODEL_ID = "TomoroAI/tomoro-colqwen3-embed-4b"
@@ -252,42 +250,8 @@ model = AutoModel.from_pretrained(
     device_map=DEVICE,
 ).eval()
 
-queries = ["Retrieve the football video", "Find the basketball clip"]
-videos = ["sample_videos/football.mp4", "sample_videos/basketball.mp4"]
-
-
-def _pad_video_sequences(batch: Dict[str, Any], video_paths: list[str]) -> Dict[str, Any]:
-    """Recompute and pad flattened video patches so they align with grid metadata."""
-    pixel_values = batch.get("pixel_values_videos")
-    grid = batch.get("video_grid_thw")
-    if not isinstance(pixel_values, torch.Tensor) or not isinstance(grid, torch.Tensor):
-        return batch
-    if pixel_values.ndim == 3:
-        return batch
-
-    rebuilt = processor.video_processor(
-        videos=video_paths,
-        return_tensors="pt",
-        return_metadata=True,
-        data_format="channels_first",
-        do_convert_rgb=True,
-    )
-    seq_grid = rebuilt["video_grid_thw"]
-    flat_pixels = rebuilt["pixel_values_videos"]
-    offsets = (seq_grid[:, 0] * seq_grid[:, 1] * seq_grid[:, 2]).tolist()
-
-    sequences = []
-    cursor = 0
-    for offset in offsets:
-        next_cursor = cursor + offset
-        sequences.append(flat_pixels[cursor:next_cursor])
-        cursor = next_cursor
-
-    batch["pixel_values_videos"] = pad_sequence(sequences, batch_first=True)
-    batch["video_grid_thw"] = seq_grid
-    if "video_metadata" in batch:
-        batch["video_metadata"] = rebuilt.get("video_metadata", batch["video_metadata"])
-    return batch
+queries = ["Retrieve the football video", "Find the basketball clip", "Find the swimming clip", "Find the wrestling clip"]
+videos = ["/root/sample_videos/football.mp4", "/root/sample_videos/basketball.mp4", "/root/sample_videos/swimming.mp4",  "/root/sample_videos/wrestling.mp4"]
 
 
 def encode_queries(texts):
@@ -301,12 +265,12 @@ def encode_videos(paths):
     vids = [str(Path(p).expanduser()) for p in paths]
     feats = processor(
         videos=vids,
-        return_tensors="pt",
         padding="longest",
+        return_tensors=None,  # keep metadata as Python objects until we drop it
         videos_kwargs={"return_metadata": True},
     )
-    feats = _pad_video_sequences(feats, vids)
     feats.pop("video_metadata", None)  # drop metadata before forwarding to the model
+    feats = feats.convert_to_tensors(tensor_type="pt")
     feats = {k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v for k, v in feats.items()}
     with torch.inference_mode():
         out = model(**feats)