CR feedback

src/autora/doc/pipelines/data.py

@@ -6,7 +6,7 @@ def load_data(data_file: str) -> Tuple[List[str], List[str]]:
 
     with jsonlines.open(data_file) as reader:
         items = [item for item in reader]
-        inputs = [f"{item['instruction']}" for item in items]
+        inputs = [item["instruction"] for item in items]
         labels = [item["output"] for item in items]
         return inputs, labels
 

src/autora/doc/pipelines/train.py

@@ -26,14 +26,15 @@ def gen() -> Iterable[Dict[str, str]]:
 
 def fine_tune(base_model: str, new_model_name: str, dataset: Dataset) -> None:
     cuda_available = torch.cuda.is_available()
+    config = {}
+
     # train using 4 bit quantization for lower GPU memory usage
-    kwargs = (
-        {"device_map": "auto", "quantization_config": get_quantization_config()} if cuda_available else {}
-    )
+    if cuda_available:
+        config.update({"device_map": "auto", "quantization_config": get_quantization_config()})
 
     model = AutoModelForCausalLM.from_pretrained(
         base_model,
-        **kwargs,
+        **config,
     )
     model.config.use_cache = False
     model.config.pretraining_tp = 1