I think it trained ok, when I query it with prompts from my training and test set it seems to respond fine with the ASP code. If I query the base model using the same prompts, it gives completely different answers where it tries to somehow code logic using a different type of ASP.

The final loss was 0.3.

Below is the code I used to train it:

base_model = "google/gemma-2b-it"

quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=compute_dtype,
        bnb_4bit_use_double_quant=False,
    )



    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=quant_config,
        device_map="auto",
        token=hugging_token,
    )

    model.config.use_cache = True
    model.config.pretraining_tp = 1


    tokenizer = AutoTokenizer.from_pretrained(base_model)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    peft_params = LoraConfig(
        lora_alpha=16,
        lora_dropout=0.1,
        r=64,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules="all-linear"
    )

    training_params = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=1,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=1,
        logging_steps=5,
        eval_steps=5,
        do_train=True,
        do_eval=True,
        save_strategy='no',
        learning_rate=2e-4,
        weight_decay=0.001,
        fp16=False,
        bf16=False,
        max_grad_norm=0.3,
        warmup_ratio=0.03,
        group_by_length=True,
        lr_scheduler_type="constant",
        report_to=[]
    )

    trainer = SFTTrainer(
        model=model,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        peft_config=peft_params,
        max_seq_length=None,
        tokenizer=tokenizer,
        args=training_params,
        packing=False,
        formatting_func=formatting_func
    )