Added more detailed logging.

Author: adivekar-utexas

src/fmcore/algorithm/huggingface/transformers.py

@@ -1,4 +1,5 @@
 import gc
+import os
 import time
 from abc import ABC
 from collections import OrderedDict
@@ -35,6 +36,8 @@
 from fmcore.constants import MLType
 
 with optional_dependency("torch", "sentencepiece", "transformers", "tokenizers", "huggingface_hub"):
+    os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "300"
+
     import huggingface_hub
     import torch
     from torch import Tensor
@@ -70,14 +73,6 @@
     )
 
     from fmcore.framework import Dataset
-    from fmcore.framework.dl.torch import (
-        Loss,
-        Optimizer,
-        PyTorch,
-        PyTorchBaseModel,
-        PyTorchClassifierMixin,
-        PyTorchMultiLabelClassifierMixin,
-    )
     from fmcore.framework._task.text_generation import (
         GENERATED_TEXTS_COL,
         GenerationOutputScoresFormat,
@@ -86,6 +81,14 @@
         TextGenerationParams,
         TextGenerationParamsMapper,
     )
+    from fmcore.framework.dl.torch import (
+        Loss,
+        Optimizer,
+        PyTorch,
+        PyTorchBaseModel,
+        PyTorchClassifierMixin,
+        PyTorchMultiLabelClassifierMixin,
+    )
 
     def mapping_to_auto_model_classes(mapping_names: Union[List, Dict, OrderedDict]) -> Dict[str, str]:
         if isinstance(mapping_names, (dict, OrderedDict)):

src/fmcore/algorithm/vllm.py

@@ -29,6 +29,7 @@ class Hyperparameters(GenerativeLM.Hyperparameters):
             gpu_memory_utilization: confloat(gt=0.0, le=1.0) = 0.95
             max_model_len: conint(ge=1)
             generation_params: Union[TextGenerationParams, Dict, str]
+            api_key: Optional[str] = None
 
             @model_validator(mode="before")
             @classmethod
@@ -46,14 +47,25 @@ def set_params(cls, params: Dict) -> Dict:
                     params,
                     param="max_model_len",
                     alias=[
+                        "max_length",
                         "max_len",
-                        "max_model_len",
                         "max_sequence_length",
                         "max_sequence_len",
                         "max_input_length",
                         "max_input_len",
+                        "max_model_length",
+                        "max_model_len",
+                    ],
+                )
+                set_param_from_alias(
+                    params,
+                    param="api_key",
+                    alias=[
+                        "token",
+                        "api_token",
                     ],
                 )
+
                 params["generation_params"] = TextGenerationParamsMapper.of(
                     params["generation_params"]
                 ).initialize()
@@ -74,10 +86,11 @@ def initialize(self, model_dir: Optional[FileMetadata] = None):
                 gpu_memory_utilization=self.hyperparams.gpu_memory_utilization,
                 max_model_len=self.hyperparams.max_model_len,
             )
-
+            kwargs["hf_overrides"]: Dict = dict()
             if self.cache_dir is not None:
                 kwargs["download_dir"] = self.cache_dir.path
-
+            if self.hyperparams.api_key is not None:
+                kwargs["hf_overrides"]["api_key"] = self.hyperparams.api_key
             print(f"Initializing vllm with kwargs: {kwargs}")
             self.llm = LLM(**kwargs)
 
@@ -103,6 +116,7 @@ def predict_step(self, batch: Prompts, **kwargs) -> Dict:
             outputs = self.llm.generate(
                 prompts,
                 sampling_params=sampling_params,
+                use_tqdm=False,
             )
 
             result = {GENERATED_TEXTS_COL: [output.outputs[0].text for output in outputs]}

src/fmcore/framework/_algorithm.py

@@ -230,6 +230,7 @@ def create_hyperparams(cls, hyperparams: Optional[Dict] = None) -> Hyperparamete
     @classmethod
     def convert_params(cls, params: Dict) -> Dict:
         ## Convert and validate parameters for the algorithm
+        # print(f'params for {cls.class_name}=\n{params}')
         cls.set_default_param_values(params)
         ## This allows us to create a new Algorithm instance without specifying `hyperparams`.
         ## If it is specified, we will pick cls.Hyperparameters, which can be overridden by the subclass.

src/fmcore/framework/_chain/Chain.py

@@ -291,8 +291,8 @@ def run(
         background: bool = False,
         tracker: Optional[Union[Tracker, Dict, str]] = None,
         notifier: Optional[Union[Notifier, Dict, str]] = None,
-        store_step_inputs: bool = False,
-        store_step_outputs: bool = False,
+        store_step_inputs: bool = True,
+        store_step_outputs: bool = True,
         after: Optional[ChainExecution] = None,
         after_wait: conint(ge=0) = 15,
         step_wait: confloat(ge=0.0) = 0.0,

src/fmcore/framework/_dataset.py

@@ -85,6 +85,10 @@ def concat(
     @model_validator(mode="before")
     @classmethod
     def _set_dataset_params(cls, params: Dict) -> Dict:
+        if "data_schema" not in params:
+            raise ValueError(
+                f"Cannot create instance of class '{cls.class_name}' without passing `data_schema` parameter."
+            )
         data_schema: Union[Schema, MLTypeSchema] = params["data_schema"]
         if isinstance(data_schema, dict):
             ## We need to infer the schema:

src/fmcore/framework/_predictions.py

@@ -36,8 +36,8 @@
     MLTypeSchema,
     TaskOrStr,
 )
-from fmcore.framework._task_mixins import InputOutputDataMixin, SchemaValidationError
 from fmcore.framework._dataset import Dataset
+from fmcore.framework._task_mixins import InputOutputDataMixin, SchemaValidationError
 
 Predictions = "Predictions"
 Visualization = "Visualization"
@@ -69,6 +69,11 @@ def _pre_registration_hook(cls):
     @model_validator(mode="before")
     @classmethod
     def _set_predictions_params(cls, params: Dict) -> Dict:
+        if "data_schema" not in params:
+            raise ValueError(
+                f"Cannot create instance of class '{cls.class_name}' without passing `data_schema` parameter."
+            )
+
         params["data_schema"]: Schema = Schema.of(params["data_schema"], schema_template=cls.schema_template)
         # data_schema: Union[Schema, MLTypeSchema] = params['data_schema']
         # if isinstance(data_schema, dict):

src/fmcore/framework/_task/text_generation.py

@@ -752,12 +752,12 @@ def _create_predictions(self, batch: Prompts, predictions: Dict, **kwargs) -> Ne
 
 
 class LanguageModelTaskMixin(Algorithm, ABC):
-    lm: Optional[Union[GenerativeLM, Any]] = None
+    lm: Optional[Any] = None
     icl_dataset: Optional[Dataset] = None
     icl_sampler: Optional[ICLSampler] = None  ## Will be not-None when icl_dataset is not-None.
 
     class Hyperparameters(Algorithm.Hyperparameters):
-        lm: Optional[Dict]  ## Params for llm
+        lm: Optional[Dict] = None  ## Params for llm
         batch_size: Optional[conint(ge=1)] = 1  ## By default, predict 1 row at a time.
         prompt_template: constr(min_length=1)
         icl_template: Optional[constr(min_length=1)] = None