main.py

import wandb
import os
import asyncio
import random
import re
import pandas as pd
import numpy as np
import tinker
import json
from tinker import types
from sklearn.model_selection import train_test_split
from tinker_cookbook.supervised.common import compute_mean_nll
from tinker_cookbook import renderers
from dotenv import load_dotenv
import warnings
warnings.filterwarnings('ignore')

load_dotenv()

API_KEY = os.environ.get("API_KEY", "")
WANDB_API_KEY = os.environ.get("WANDB_API_KEY", '')
BASE_MODEL = "Qwen/Qwen3-30B-A3B"
DATASET_PATH = "final_merged_table_cut.json"
TEXT_COL = "text"            
ATTR_COLS = ["sentiment", "country", "gender"]
BATCH_SIZE = 4
LEARNING_RATE = 1e-4
MAX_TOKENS_INFERENCE = 100
RANDOM_SEED = 42

random.seed(RANDOM_SEED)
np.random.seed(RANDOM_SEED)


    print(f"Loading dataset from {path}...")
    df = pd.read_json(path)
    
    # Validate required columns
    required_cols = [TEXT_COL] + ATTR_COLS
    missing_cols = [col for col in required_cols if col not in df.columns]
    if missing_cols:
        raise ValueError(f"Missing required columns: {missing_cols}")
    
    # Clean data: remove rows with NaN
    df_clean = df.dropna(subset=required_cols).copy()
    if len(df) != len(df_clean):
        print(f"Removed {len(df) - len(df_clean)} rows with missing values")
    
    # 1. Split data: 22% for testing, 78% for training
    train_pool, test_df = train_test_split(
        df_clean, 
        test_size=test_size, 
        random_state=RANDOM_SEED
    )
    
    print(f"Total Data Stats:")
    print(f"  Training Pool ({100*(1-test_size):.0f}%): {len(train_pool)} rows")
    print(f"  Test Set ({100*test_size:.0f}%):      {len(test_df)} rows")
    
    return train_pool, test_df

def format_prompt(text, target_attr=None):
    if target_attr:
        # model b
        return f"Analyze the text and extract the {target_attr}.\n\nText: {text}\n\nResponse:\n"
    else:
        # model a
        return f"Analyze the text and extract attributes.\n\nText: {text}\n\nAttributes:\n"

def export_error_analysis(sampling_client_A, sampling_client_B, tokenizer, test_df, output_file="error_analysis.csv"):
    # export failures for error analysis
    print(f"Starting Error Analysis Export...")
    
    eval_df = test_df.head(200).copy()
    records = []
    
    renderer = renderers.get_renderer('qwen3', tokenizer)
    stop_sequences = renderer.get_stop_sequences()

    for i, (idx, row) in enumerate(eval_df.iterrows()):
        input_text = row[TEXT_COL]
        
        # 1. Get Model A Predictions (One inference)
        prompt_A = format_prompt(input_text, target_attr=None)
        input_A = types.ModelInput.from_ints(tokenizer.encode(prompt_A, add_special_tokens=True))
        res_A = sampling_client_A.sample(prompt=input_A, num_samples=1, 
                                        sampling_params=types.SamplingParams(max_tokens=MAX_TOKENS_INFERENCE, temperature=0.01, stop=stop_sequences)).result()
        gen_A = tokenizer.decode(res_A.sequences[0].tokens, skip_special_tokens=True)
        preds_A = parse_attributes(gen_A, ATTR_COLS)

        # 2. Get Model B Predictions (Three inferences)
        preds_B = {}
        for col in ATTR_COLS:
            prompt_B = format_prompt(input_text, target_attr=col)
            input_B = types.ModelInput.from_ints(tokenizer.encode(prompt_B, add_special_tokens=True))
            res_B = sampling_client_B.sample(prompt=input_B, num_samples=1, 
                                            sampling_params=types.SamplingParams(max_tokens=20, temperature=0.01, stop=stop_sequences)).result()
            gen_B = tokenizer.decode(res_B.sequences[0].tokens, skip_special_tokens=True)
            parsed_B = parse_attributes(gen_B, [col])
            preds_B[col] = parsed_B.get(col, "")

        # 3. Create record for CSV
        record = {
            "text": input_text,
            "true_sentiment": row["sentiment"],
            "true_country": row["country"],
            "true_gender": row["gender"],
            "model_A_sentiment": preds_A.get("sentiment", ""),
            "model_A_country": preds_A.get("country", ""),
            "model_A_gender": preds_A.get("gender", ""),
            "model_B_sentiment": preds_B.get("sentiment", ""),
            "model_B_country": preds_B.get("country", ""),
            "model_B_gender": preds_B.get("gender", ""),
        }
        
        # Add "Is Correct" flags for easier filtering in Excel
        for col in ATTR_COLS:
            record[f"A_{col}_correct"] = str(row[col]).lower() == str(preds_A.get(col, "")).lower()
            record[f"B_{col}_correct"] = str(row[col]).lower() == str(preds_B.get(col, "")).lower()
            
        records.append(record)
        if (i+1) % 10 == 0:
            print(f"Processed {i+1}/200 samples...")

    # Export to CSV
    error_df = pd.DataFrame(records)
    error_df.to_csv(output_file, index=False)
    print(f"✅ Error analysis file saved to: {output_file}")

def process_batch(df_batch, tokenizer, mask_mode="none"):
    """
    Process a batch of data for training.
    """
    processed_examples = []
    
    for _, row in df_batch.iterrows():
        # 1. Determine Masking & Target Attribute FIRST
        target_attr = None
        mask_flags = [False] * len(ATTR_COLS)

        if mask_mode == "random_single":
            # Pick one random attribute to KEEP
            keep_idx = random.randint(0, len(ATTR_COLS) - 1)
            target_attr = ATTR_COLS[keep_idx]
            
            # Mask everything...
            mask_flags = [True] * len(ATTR_COLS)
            # ...except the target
            mask_flags[keep_idx] = False
            
        # 2. Generate Prompt based on decision
        # If target_attr is set (Model B), we get specific prompt.
        # If None (Model A), we get generic prompt.
        prompt_text = format_prompt(row[TEXT_COL], target_attr)
        
        prompt_tokens = tokenizer.encode(prompt_text, add_special_tokens=True)
        full_tokens = list(prompt_tokens)
        full_weights = [0.0] * len(prompt_tokens)
        
        # 3. Build Token Sequence for attributes
        for i, col in enumerate(ATTR_COLS):
            # Skip masked attributes
            if mask_flags[i]:
                continue

            val = str(row[col])
            line = f"{col}: {val}\n"
            
            # Encode without special tokens to avoid adding BOS/EOS between attributes
            line_tokens = tokenizer.encode(line, add_special_tokens=False)
            
            weight = 1.0
            line_weights = [weight] * len(line_tokens)
            
            full_tokens.extend(line_tokens)
            full_weights.extend(line_weights)

        eos_id = tokenizer.eos_token_id
        if eos_id is None:
            eos_id = 151645 
            
        full_tokens.append(eos_id)
        full_weights.append(1.0) 
        
        input_tokens_model = full_tokens[:-1]
        target_tokens = full_tokens[1:]
        weights_shifted = full_weights[1:]
        
        datum = types.Datum(
            model_input=types.ModelInput.from_ints(input_tokens_model),
            loss_fn_inputs={"weights": weights_shifted, "target_tokens": target_tokens}
        )

        processed_examples.append(datum)
    
    return processed_examples

def train_model(service_client, train_df, mask_mode, model_name, epochs=1):
    print(f"\n--- Training {model_name} ---")
    print(f"    Mode: {mask_mode}")
    print(f"    Data Size: {len(train_df)} rows")
    print(f"    Epochs: {epochs}")
    
    checkpoint_file = f"{model_name}_checkpoint.pt"
    start_epoch = 0
    resume_path = None

    if os.path.exists(checkpoint_file):
        try:
            with open(checkpoint_file, 'r') as f:
                state = json.load(f)
                resume_path = state['checkpoint_path']
                start_epoch = state.get("next_epoch", 0)
            if resume_path:
                print(f"found checkpoint file, starting from epoch {start_epoch}")
        except Exception as e:
            print(f"Error loading checkpoint file: {e}")

    training_client = service_client.create_lora_training_client(
        base_model=BASE_MODEL,
        rank=32,
    )
    tokenizer = training_client.get_tokenizer()
    
    if resume_path:
        print("   Loading optimizer state and weights from cloud...")
        load_future = training_client.load_state(resume_path)
        if hasattr(load_future, "result"):
            load_future.result()
        print("   State loaded successfully.")

    num_batches = (len(train_df) * epochs) // BATCH_SIZE
    if num_batches == 0:
        print("Warning: Not enough data for a single batch!")
        return training_client.save_weights_and_get_sampling_client(name=model_name), tokenizer
    
    if start_epoch >= epochs:
        print("Training already completed based on checkpoint.")
        return training_client.save_weights_and_get_sampling_client(name=model_name), tokenizer

    # --- Training Loop ---
    for epoch in range(start_epoch, epochs):
        print(f"\nEpoch {epoch + 1}/{epochs}")
        train_df_shuffled = train_df.sample(frac=1, random_state=RANDOM_SEED + epoch).reset_index(drop=True)
        
        for i in range(0, len(train_df_shuffled), BATCH_SIZE):
            batch_idx = i // BATCH_SIZE + epoch * (len(train_df_shuffled) // BATCH_SIZE)
            batch_df = train_df_shuffled.iloc[i:i + BATCH_SIZE]
            
            if len(batch_df) < 1: continue
            
            data = process_batch(batch_df, tokenizer, mask_mode=mask_mode)
            
            fwd_future = training_client.forward_backward(data, "cross_entropy")
            optim_future = training_client.optim_step(
                types.AdamParams(learning_rate=LEARNING_RATE)
            )
            
            fwd_bwd_result = fwd_future.result()
            optim_future.result()

            logprobs = [x["logprobs"] for x in fwd_bwd_result.loss_fn_outputs]
            weights = [datum.loss_fn_inputs["weights"] for datum in data]
            train_nll = compute_mean_nll(logprobs, weights)

            wandb.log({
                f"{model_name}/train/loss": train_nll,
                "batch_idx": batch_idx,
                "epoch": epoch
            })
            
            if batch_idx % 10 == 0:
                print(f"[{model_name}] Batch {batch_idx}/{num_batches} processed")

        # --- Save State ---
        print(f"Saving full state for Epoch {epoch+1}...")
        try:
            ckpt_name = f"{model_name}_epoch_{epoch+1}"
            saved_path = training_client.save_state(name=ckpt_name).result().path
            state = {"checkpoint_path": saved_path, "next_epoch": epoch + 1}
            with open(checkpoint_file, 'w') as f:
                json.dump(state, f)
            print(f"State saved to: {saved_path}")
        except Exception as e:
            print(f"Failed to save state: {e}")

    print(f"[{model_name}] Training finished.")
    sampling_client = training_client.save_weights_and_get_sampling_client(name=model_name)
    return sampling_client, tokenizer

def parse_attributes(text, attr_cols):
    result = {}
    lines = text.strip().split('\n')
    for line in lines:
        line = line.strip()
        for attr in attr_cols:
            if line.lower().startswith(attr.lower() + ":"):
                value = line.split(':', 1)[1].strip()
                result[attr] = value
                break
    for attr in attr_cols:
        if attr not in result:
            result[attr] = ""
    return result

def normalize_attribute(value):
    if not isinstance(value, str):
        value = str(value)
    return value.strip().lower()

def evaluate_model(sampling_client, tokenizer, test_df, model_name, max_samples=None):
    print(f"\n--- Evaluating {model_name} ---")
    
    if max_samples and max_samples < len(test_df):
        eval_df = test_df.head(max_samples)
        print(f"Evaluating on {len(eval_df)} samples (subsampled)")
    else:
        eval_df = test_df
        print(f"Evaluating on {len(eval_df)} samples")
    
    correct_counts = {col: 0 for col in ATTR_COLS}
    total_samples = len(eval_df)
    
    renderer = renderers.get_renderer('qwen3', tokenizer)
    stop_sequences = renderer.get_stop_sequences()

    for i, (idx, row) in enumerate(eval_df.iterrows()):
        row_predictions = {}
        
        if model_name == "Model_B_Masked":
            # Model B: Loop through every attribute (3 inferences per row)
            for col in ATTR_COLS:
                prompt_text = format_prompt(row[TEXT_COL], target_attr=col)
                prompt_input = types.ModelInput.from_ints(
                    tokenizer.encode(prompt_text, add_special_tokens=True)
                )
                
                try:
                    response_future = sampling_client.sample(
                        prompt=prompt_input,
                        num_samples=1,
                        sampling_params=types.SamplingParams(
                            max_tokens=20, # Short generation needed for single item
                            temperature=0.01,
                            stop=stop_sequences
                        )
                    )
                    result = response_future.result()
                    gen_text = tokenizer.decode(result.sequences[0].tokens, skip_special_tokens=True)
                    
                    # Parse specifically for this column
                    parsed = parse_attributes(gen_text, [col])
                    row_predictions[col] = parsed.get(col, "")
                except Exception as e:
                    print(f"Error inference Model B on {col}: {e}")
                    row_predictions[col] = ""
                    
        else:
            # Model A: Run ONCE with generic prompt (1 inference per row)
            prompt_text = format_prompt(row[TEXT_COL], target_attr=None)
            prompt_input = types.ModelInput.from_ints(
                tokenizer.encode(prompt_text, add_special_tokens=True)
            )
            
            try:
                response_future = sampling_client.sample(
                    prompt=prompt_input,
                    num_samples=1,
                    sampling_params=types.SamplingParams(
                        max_tokens=MAX_TOKENS_INFERENCE,
                        temperature=0.01,
                        stop=stop_sequences
                    )
                )
                result = response_future.result()
                gen_text = tokenizer.decode(result.sequences[0].tokens, skip_special_tokens=True)
                
                # Parse all attributes at once
                row_predictions = parse_attributes(gen_text, ATTR_COLS)
            except Exception as e:
                print(f"Error inference Model A: {e}")
                row_predictions = {}

        # Scoring
        for col in ATTR_COLS:
            truth = normalize_attribute(row[col])
            pred = normalize_attribute(row_predictions.get(col, ""))
            if truth == pred:
                correct_counts[col] += 1
        
        # heartbeat output
        if i % 50 == 0: 
            print(f"  Sample {i+1}/{total_samples}")
    
    print(f"\nResults for {model_name}:")
    results = {}
    for col in ATTR_COLS:
        acc = correct_counts[col] / total_samples if total_samples > 0 else 0
        results[col] = acc
        print(f"  {col} Accuracy: {acc:.2%} ({correct_counts[col]}/{total_samples})")
    
    avg_acc = sum(results.values()) / len(results) if results else 0
    print(f"  Average Accuracy: {avg_acc:.2%}")
    
    return results

def main():
    wandb.init(
        project="supervised-llm-finetuning",
        name="qwen-lora-run-1"
    )

    if not API_KEY:
        raise ValueError("API_KEY not found.")
    
    service_client = tinker.ServiceClient(api_key=API_KEY)
    
    train_pool, test_df = load_and_split_data(DATASET_PATH)
    
    train_df_model_A = train_pool.sample(frac=1/3, random_state=RANDOM_SEED)
    train_df_model_B = train_pool.copy()
    
    print("\n" + "="*50)
    print("TRAINING CONFIGURATION:")
    print("="*50)
    print(f"Model A (Baseline): {len(train_df_model_A)} rows")
    print(f"Model B (Masked):   {len(train_df_model_B)} rows")
    print("="*50)
    
    # 2. Train Model A
    sampler_A, tokenizer = train_model(
        service_client, 
        pd.DataFrame(), 
        mask_mode="none", 
        model_name="Model_A_Baseline",
        epochs=1
    )
    
    # 3. Train Model B
    sampler_B, _ = train_model(
        service_client, 
        pd.DataFrame(), 
        mask_mode="random_single", 
        model_name="Model_B_Masked",
        epochs=1
    )

    print("\n" + "="*50)
    print("RUNNING 200-SAMPLE ERROR ANALYSIS EXPORT")
    print("="*50)
    
    export_error_analysis(
        sampler_A, 
        sampler_B, 
        tokenizer, 
        test_df, 
        output_file="error_analysis_results.csv"
    )
    
    # 4. Evaluation
    print("\n" + "="*50)
    print("FINAL EVALUATION")
    print("="*50)
        results_A = evaluate_model(
        sampler_A, 
        tokenizer, 
        test_df, 
        "Model_A_Baseline",
        max_samples=1000  # number of samples for eval
    )
    
    results_B = evaluate_model(
        sampler_B, 
        tokenizer, 
        test_df, 
        "Model_B_Masked",
        max_samples=1000
    )
    
    # 5. Summary
    print("\n" + "="*50)
    print("SUMMARY COMPARISON")
    print("="*50)
    print(f"{'Attribute':<15} {'Model A':<12} {'Model B':<12} {'Difference':<12}")
    print("-"*50)
    
    for attr in ATTR_COLS:
        acc_A = results_A.get(attr, 0)
        acc_B = results_B.get(attr, 0)
        diff = acc_B - acc_A
        diff_sign = "+" if diff > 0 else ""
        print(f"{attr:<15} {acc_A:.2%}{'':<4} {acc_B:.2%}{'':<4} {diff_sign}{diff:.2%}")
    
    avg_A = sum(results_A.values()) / len(results_A) if results_A else 0
    avg_B = sum(results_B.values()) / len(results_B) if results_B else 0
    avg_diff = avg_B - avg_A
    diff_sign = "+" if avg_diff > 0 else ""
    print("-"*50)
    print(f"{'Average':<15} {avg_A:.2%}{'':<4} {avg_B:.2%}{'':<4} {diff_sign}{avg_diff:.2%}")

if __name__ == "__main__":
    main()