feat: Add testing mode support for Arctic encoder and subset selection

- Introduce testing_mode flag in EncoderConfig for Arctic encoder
- Enable direct HuggingFace model download in testing mode
- Modify subset_datasets to pass testing_mode to encoder configuration
- Add warning for local model download in testing mode
- Improve flexibility for model loading during testing scenarios

Signed-off-by: eshwarprasadS <eshwarprasad.s01@gmail.com>

Author: eshwarprasadS

src/instructlab/sdg/encoders/arctic_encoder.py

@@ -43,7 +43,7 @@ class ModelConfig(TypedDict):
     }
 }
 
-
+# pylint: disable=too-many-instance-attributes
 @dataclass
 class EncoderConfig:
     model_name: str
@@ -53,6 +53,7 @@ class EncoderConfig:
     batch_size: int
     use_default_instruction: bool
     use_fp16: bool
+    testing_mode: bool = False
 
 
 class ArcticEmbedEncoder:
@@ -62,6 +63,7 @@ def __init__(
         device: Optional[torch.device] = None,
         use_fp16: bool = False,
         use_default_instruction: bool = True,
+        testing_mode: bool = False,
     ) -> None:
         """Initialize the Arctic encoder."""
         if model_name not in MODEL_CONFIGS:
@@ -82,6 +84,7 @@ def __init__(
             batch_size=batch_size,
             use_default_instruction=use_default_instruction,
             use_fp16=use_fp16,
+            testing_mode=testing_mode,
         )
 
         self._initialize_model()
@@ -93,11 +96,24 @@ def _initialize_model(self) -> None:
             home_dir, ".cache", "instructlab", "models", self.cfg.model_name
         )
 
-        if not os.path.exists(model_path):
-            raise ValueError(
-                f"Model not found in available models: {self.cfg.model_name}\n"
-                "Please run `ilab model download` and download the necessary model"
+        # In testing mode, allow direct download from HuggingFace
+        if hasattr(self.cfg, "testing_mode") and self.cfg.testing_mode:
+            logger.warning(
+                f"Model not found locally at {model_path}. "
+                "Testing mode enabled - downloading from HuggingFace..."
+            )
+            self.tokenizer = AutoTokenizer.from_pretrained(self.cfg.model_name)
+            self.model = AutoModel.from_pretrained(
+                self.cfg.model_name,
+                add_pooling_layer=False,
+                trust_remote_code=True,
             )
+        else:
+            if not os.path.exists(model_path):
+                raise ValueError(
+                    f"Model not found in available models: {self.cfg.model_name}\n"
+                    "Please run `ilab model download` and download the necessary model"
+                )
 
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
         self.model = AutoModel.from_pretrained(

src/instructlab/sdg/subset_selection.py

@@ -90,6 +90,7 @@ class EncoderConfig:
     encoder_model: str = field(
         default="Snowflake/snowflake-arctic-embed-l-v2.0", metadata={"advanced": True}
     )
+    testing_mode: bool = False
 
 
 @dataclass
@@ -178,7 +179,10 @@ def __init__(self, config: ProcessingConfig, encoder_cls):
             encoder_cls: The encoder class to use for generating embeddings.
         """
         self.config = config
-        self.encoder = encoder_cls(model_name=config.encoder.encoder_model)
+        self.encoder = encoder_cls(
+            model_name=config.encoder.encoder_model,
+            testing_mode=config.encoder.testing_mode,
+        )
         self.env = Environment(loader=BaseLoader())
         self.templates = {
             k: self.env.from_string(v) for k, v in config.template.templates.items()
@@ -880,7 +884,7 @@ def subset_datasets(
 
     # Create configuration groups
     basic_config = BasicConfig()
-    encoder_config = EncoderConfig()
+    encoder_config = EncoderConfig(testing_mode=testing_mode)
     template_config = TemplateConfig()
     system_config = SystemConfig(testing_mode=testing_mode)