# config.py

import torch
from transformers import BitsAndBytesConfig
from peft import LoraConfig

# --------------------
# General
# --------------------
MODEL_NAME = "stabilityai/stable-code-3b"
OUTPUT_DIR = "finetuned_model"
MAX_LENGTH = 512
SEED = 42

# --------------------
# Dataset
# --------------------
MBPP_DATASET = ("google-research-datasets/mbpp", "full")
LOCAL_DATA_PATH = "training_data.jsonl"
TEST_SIZE = 0.1

# --------------------
# Quantization
# --------------------
BNB_CONFIG = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# --------------------
# LoRA
# --------------------
LORA_CONFIG = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    task_type="CAUSAL_LM"
)

# --------------------
# Training
# --------------------
TRAINING_ARGS = {
    "per_device_train_batch_size": 1,
    "gradient_accumulation_steps": 4,
    "num_train_epochs": 1,
    "logging_steps": 10,
    "eval_strategy": "steps",
    "eval_steps": 50,
    "fp16": True,
    "report_to": "none",
    "save_strategy": "epoch"
}

# --------------------
# Generation
# --------------------
GENERATION_CONFIG = {
    "max_new_tokens": 300,
    "temperature": 0.7,
    "do_sample": True
}