diff --git a/src/autora/doc/pipelines/data.py b/src/autora/doc/pipelines/data.py index fc916a7..e067297 100644 --- a/src/autora/doc/pipelines/data.py +++ b/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 diff --git a/src/autora/doc/pipelines/train.py b/src/autora/doc/pipelines/train.py index ad81518..b001fa3 100644 --- a/src/autora/doc/pipelines/train.py +++ b/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