update code

Author: ParisZhang11

examples/post_training_llava/README.md

@@ -20,7 +20,7 @@ source .venv/bin/activate
 
 ## Example
 
-Please update the fields `annotation_path` and `media_path` in `configs/sft.toml` to your custom dataset. `media_path` can be left as empty (`""`) if the paths in your annotation are absolute paths. 
+Please update the fields `annotation_path` and `media_path` in `configs/sft.toml` to your custom dataset. `media_path` can be left as empty (`""`) if the paths in your annotation are absolute paths.
 
 Here is one example of downloading the [Llava-Instruct-150K] (https://huggingface.co/datasets/liuhaotian/LLaVA-Instruct-150K) dataset and [COCO] (https://cocodataset.org/#home) images:
 

examples/post_training_llava/scripts/custom_sft.py

@@ -18,6 +18,7 @@
 import argparse
 import json
 import os
+import re
 from pathlib import Path
 
 import cosmos_rl.launcher.worker_entry
@@ -30,6 +31,7 @@
 from cosmos_reason1_utils.text import create_conversation
 from cosmos_reason1_utils.vision import VisionConfig
 
+
 class CustomDatasetConfig(pydantic.BaseModel):
     annotation_path: str = pydantic.Field()
     """Dataset annotation path."""
@@ -85,7 +87,6 @@ def __getitem__(self, idx: int) -> list[dict]:
                 videos = [os.path.join(self.media_path, vid) for vid in videos]
 
         # Remove image and video tags from user prompt
-        import re
         user_prompt = re.sub(r"(\n)?</?(image|video)>(\n)?", "", user_prompt)
 
         conversations = create_conversation(
@@ -97,6 +98,7 @@ def __getitem__(self, idx: int) -> list[dict]:
         )
         return conversations
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(