# ============================================================================
# === CONFIGURATION - ALL SETTINGS IN ONE PLACE ===
# ============================================================================

# --- Model Configuration ---
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" # YOU CANNOT CHANGE THIS

# --- Dataset Configuration ---
#TODO: REPLACE WITH YOUR OWN PATH
MCQ_CSV_PATH = "hw5_sample_eval.csv"  # Path to CS189 MCQ sample eval dataset

# --- Training Configuration (feel free to adjust!) ---
TRAIN_BATCH_SIZE = 1
GRADIENT_ACCUMULATION_STEPS = 4
WARMUP_STEPS = 5
MAX_STEPS = 50  # or set num_train_epochs instead
LEARNING_RATE = 1e-5
WEIGHT_DECAY = 0.01
LR_SCHEDULER_TYPE = "linear"
OPTIM = "adamw_8bit"  # requires bitsandbytes
SEED = 189

# --- Evaluation Configuration ---
EVAL_MAX_NEW_TOKENS = 64  # How many tokens to generate for inference
OUTPUT_DIR = "./mcq_finetuned_model"

# === Load base model & tokenizer ===
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Ensure we have a pad token for training
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None,
)
model.resize_token_embeddings(len(tokenizer))
model.to(device)
model.eval()
print('Model loaded.')

# === MCQ helpers ===
LETTER_SET = set(list("ABCDE"))

def load_mcq_dataset(csv_path: str = MCQ_CSV_PATH):
    """Load the CS189 MCQ dataset.

    Expected columns:
        - question
        - A, B, C, D, E
        - answer (single letter A-E)
    """
    df = pd.read_csv(csv_path)
    required = ["question", "A", "B", "C", "D", "E", "answer"]
    missing = [c for c in required if c not in df.columns]
    if missing:
        raise ValueError(f"Missing required columns in MCQ CSV: {missing}")

    df = df.copy()
    df["answer"] = (
        df["answer"]
        .astype(str)
        .str.strip()
        .str.upper()
    )
    df = df[df["answer"].isin(LETTER_SET)].reset_index(drop=True)
    return df

def build_mcq_prompt(row):
    """Prompt for inference: instruction + question + options.

    The model is expected to answer with the correct letter in \\boxed{} format.
    """
    q = str(row["question"]).strip()
    options = "\n".join([
        f"A. {row['A']}",
        f"B. {row['B']}",
        f"C. {row['C']}",
        f"D. {row['D']}",
        f"E. {row['E']}",
    ])
    prompt = (
        "Choose exactly one correct option from A, B, C, D, and E.\n"
        "Return your answer inside a LaTeX box.\n\n"
        f"{q}\n\n{options}\n\nAnswer:"
    )
    return prompt

def parse_choice_from_boxed(text: str):
    """Parse an MCQ choice A–E from the model output.

    We first look for a literal '\\boxed{X}' pattern. If not found, we
    fallback to the last standalone A-E in the decoded text.
    """
    if text is None:
        return None
    # Direct \\boxed{A} ... \\boxed{E}
    m = re.search(r"\\boxed\{\s*([A-E])\s*\}", text)
    if m:
        return m.group(1)
    # Fallback: last standalone A–E
    letters = re.findall(r"\b([A-E])\b", text.upper())
    if letters:
        return letters[-1]
    return None

# === MMLU Helper Functions ===
def load_mmlu_dataset(subset: str = "machine_learning", split: str = "test"):
    """Load a subset of the MMLU dataset from Hugging Face."""
    print(f"Loading MMLU dataset (subset={subset}, split={split})...")
    ds = load_dataset("cais/mmlu", subset, split=split)
    return ds

def build_mmlu_prompt(row):
    """Prompt for inference: instruction + question + options."""
    q = str(row["question"]).strip()
    choices = row["choices"]

    options_list = []
    for i, choice in enumerate(choices):
        letter = chr(ord("A") + i)
        options_list.append(f"{letter}. {choice}")
    options_str = "\n".join(options_list)

    prompt = (
        "Choose exactly one correct option from the choices provided.\n"
        "Return your answer inside a LaTeX box.\n\n"
        f"{q}\n\n{options_str}\n\nAnswer:"
    )
    return prompt

def build_mmlu_sft_text(row, tokenizer):
    """Build properly formatted chat template text for training."""
    user_content = build_mmlu_prompt(row)

    answer_int = row["answer"]
    answer_letter = chr(ord("A") + answer_int)
    assistant_content = f"\\boxed{{{answer_letter}}}"

    messages = [
        {"role": "user", "content": user_content},
        {"role": "assistant", "content": assistant_content}
    ]

    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False
    )

# === Load MMLU Machine Learning Dataset ===
mmlu_ds = load_mmlu_dataset("machine_learning", split="test")
mmlu_text_ds = mmlu_ds.map(lambda x: {"text": build_mmlu_sft_text(x, tokenizer)})
print("Loaded MMLU ML dataset with", len(mmlu_text_ds), "rows")

# Set the training dataset - you can mix and match datasets here
train_dataset = mmlu_text_ds

<|im_start|>system
You are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
<|im_start|>user
Choose exactly one correct option from the choices provided.
Return your answer inside a LaTeX box.

Statement 1| Linear regression estimator has the smallest variance among all unbiased estimators. Statement 2| The coefficients α assigned to the classifiers assembled by AdaBoost are always non-negative.

A. True, True
B. False, False
C. True, False
D. False, True

Answer:<|im_end|>
<|im_start|>assistant
\boxed{D}<|im_end|>

def eval_mcq_accuracy(
    curr_model,
    curr_tokenizer,
    df,
    max_new_tokens: int = 64,
    return_details: bool = False,
):
    """Evaluate a model on the MCQ dataset using greedy decoding.

    If return_details=True, also return a pandas DataFrame with
    [idx, question, A, B, C, D, E, gold, decoded, parsed, correct].
    """
    curr_model.eval()
    n = len(df)
    correct = 0
    total = 0
    records = []

    for idx in range(n):
        row = df.iloc[idx]
        user_content = build_mcq_prompt(row)

        # Apply chat template for inference
        messages = [{"role": "user", "content": user_content}]
        prompt = curr_tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        inputs = curr_tokenizer(prompt, return_tensors="pt").to(device)

        with torch.no_grad():
            outputs = curr_model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
            )

        gen_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        decoded = curr_tokenizer.decode(gen_tokens, skip_special_tokens=True)

        pred = parse_choice_from_boxed(decoded)
        is_correct = (pred is not None and pred == row["answer"])
        if is_correct:
            correct += 1
        total += 1

        records.append({
            "idx": idx,
            "question": row["question"],
            "A": row["A"],
            "B": row["B"],
            "C": row["C"],
            "D": row["D"],
            "E": row["E"],
            "gold": row["answer"],
            "prompt": prompt,
            "decoded": decoded,
            "parsed": pred,
            "correct": is_correct,
        })

        if (idx + 1) % 20 == 0:
            print(f"Processed {idx + 1}/{n} questions...")

    acc = correct / max(total, 1)
    print(f"MCQ accuracy: {acc * 100:.2f}% ({correct}/{total})")

    details_df = pd.DataFrame(records)
    if return_details:
        return acc, details_df
    return acc

# === Baseline MCQ accuracy before fine-tuning ===
print("Evaluating baseline model on MCQ dataset...")
baseline_acc, baseline_details = eval_mcq_accuracy(
    model,
    tokenizer,
    mcq_df,
    max_new_tokens=EVAL_MAX_NEW_TOKENS,
    return_details=True,
)
baseline_details.head()

# === Set up SFTTrainer ===
sft_config = SFTConfig(
    dataset_text_field="text",
    per_device_train_batch_size=TRAIN_BATCH_SIZE,
    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
    warmup_steps=WARMUP_STEPS,
    max_steps=MAX_STEPS,
    learning_rate=LEARNING_RATE,
    logging_steps=1,
    optim=OPTIM,
    weight_decay=WEIGHT_DECAY,
    lr_scheduler_type=LR_SCHEDULER_TYPE,
    seed=SEED,
    report_to="none",
)

trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=train_dataset,
    eval_dataset=None,
    processing_class=tokenizer,
)

trainer