Fix yaml import

Author: lorenmt

demo.py

@@ -1,7 +1,11 @@
 import os
 import argparse
 import torch
-import ruamel_yaml as yaml
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+
 
 from model.prismer_caption import PrismerCaption
 from dataset import create_dataset, create_loader

requirements.txt

@@ -15,4 +15,5 @@ pycocotools
 geffnet
 fire
 huggingface_hub
-rich
+rich
+ruamel.yaml

train_caption.py

@@ -5,15 +5,18 @@
 # https://github.com/NVlabs/prismer/blob/main/LICENSE
 
 import argparse
-import ruamel_yaml as yaml
 import numpy as np
 import random
 import time
 import functools
 import json
 import torch
 import os
-
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
 from accelerate import Accelerator, FullyShardedDataParallelPlugin
 from model.prismer_caption import PrismerCaption
 from model.modules.utils import interpolate_pos_embed

train_classification.py

@@ -5,12 +5,15 @@
 # https://github.com/NVlabs/prismer/blob/main/LICENSE
 
 import argparse
-import ruamel_yaml as yaml
 import numpy as np
 import random
 import time
 import functools
 import torch
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
 
 from accelerate import Accelerator, FullyShardedDataParallelPlugin
 from model.prismer_caption import PrismerCaption

train_pretrain.py

@@ -5,14 +5,17 @@
 # https://github.com/NVlabs/prismer/blob/main/LICENSE
 
 import argparse
-import ruamel_yaml as yaml
 import numpy as np
 import random
 import time
 import datetime
 import functools
 import torch
-
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
 from accelerate import Accelerator, FullyShardedDataParallelPlugin
 from model.prismer_caption import PrismerCaption
 from dataset import create_dataset, create_loader

train_vqa.py

@@ -5,14 +5,17 @@
 # https://github.com/NVlabs/prismer/blob/main/LICENSE
 
 import argparse
-import ruamel_yaml as yaml
 import numpy as np
 import random
 import time
 import datetime
 import functools
 import torch
-
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
 from accelerate import Accelerator, FullyShardedDataParallelPlugin
 from model.prismer_vqa import PrismerVQA
 from model.modules.utils import interpolate_pos_embed