Add rouge-score/anthropic/openai; make deepspeed import optional

model/dist_funs.py

@@ -6,7 +6,12 @@ from lavis.common.dist_utils import is_dist_avail_and_initialized
 from model.help_funcs import pad_and_concat
 from pytorch_lightning import strategies
 from lightning_fabric.utilities.types import _PATH
-from deepspeed.runtime.data_pipeline.data_routing.helper import remove_random_ltd_state_dict
+try:
+    from deepspeed.runtime.data_pipeline.data_routing.helper import remove_random_ltd_state_dict
+    from deepspeed import DeepSpeedEngine
+    _DEEPSPEED_AVAILABLE = True
+except ImportError:
+    _DEEPSPEED_AVAILABLE = False
 
 
 '''
@@ -20,7 +25,7 @@ def module_state_dict(self, destination=None, prefix="", keep_vars=False, exclud
     if exclude_frozen_parameters:
         to_be_removed = []
         for n in sd:
-            try: 
+            try:
                 if not self.module.get_parameter(n).requires_grad:
                     to_be_removed.append(n)
             except AttributeError:
@@ -30,11 +35,13 @@ def module_state_dict(self, destination=None, prefix="", keep_vars=False, exclud
     if self.random_ltd_enabled():
         sd = remove_random_ltd_state_dict(sd)
     return sd
-from deepspeed import DeepSpeedEngine
-DeepSpeedEngine.module_state_dict = module_state_dict
+if _DEEPSPEED_AVAILABLE:
+    DeepSpeedEngine.module_state_dict = module_state_dict
 ### end overwrite ###
 
-class MyDeepSpeedStrategy(strategies.DeepSpeedStrategy):
+_DeepSpeedStrategyBase = strategies.DeepSpeedStrategy if _DEEPSPEED_AVAILABLE else object
+
+class MyDeepSpeedStrategy(_DeepSpeedStrategyBase):
     def save_checkpoint_v1(
         self, checkpoint: Dict[str, Any], filepath: _PATH, storage_options: Optional[Any] = None
     ):

requirements.txt

@@ -48,6 +48,11 @@ regex==2024.11.6
 web.py==0.62
 cheroot==10.0.1
 
+# other runtime deps imported by model/ and evals/
+rouge-score==0.1.2
+anthropic==0.40.0
+openai==2.7.2
+
 # spacy's runtime deps (spacy itself is installed --no-deps in the Dockerfile
 # because its typer<0.10 pin conflicts with gradio's typer>=0.12)
 catalogue==2.0.10