Untitled

import os
import torch
import pandas as pd
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    set_seed,
)
from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
from trl import SFTTrainer

MODEL_NAME = "mistralai/Mistral-7B-v0.1"
OUTPUT_DIR = "./results_mistral_resume"
LORA_R = 16
LORA_ALPHA = 32
LORA_DROPOUT = 0.05
MAX_SEQ_LENGTH = 2048
LEARNING_RATE = 2e-4
BATCH_SIZE = 4
GRADIENT_ACCUMULATION_STEPS = 4 
NUM_TRAIN_EPOCHS = 3
CSV_FILE_PATH = "resume_data.csv"

set_seed(42)


def load_and_format_data(csv_path):
    """
    Loads the CSV, formats the data into the instruction-response template 
    required for Mistral-7B, and returns a Hugging Face Dataset.
    """
    print(f"Loading data from {csv_path}...")
    try:
        df = pd.read_csv(csv_path)
    except FileNotFoundError:
        print(f"Error: {csv_path} not found. Please create it first.")
        df = pd.DataFrame({
            'category': ["Data Science"],
            'resume': [
                "Education Details May 2013 to May 2017 B.E   UIT-RGPV Data Scientist Data Scientist - Matelabs Skill Details Python- Exprience - Less than 1 year months Statsmodels- Exprience - 12 months AWS- Exprience - Less than 1 year months Machine learning- Exprience - Less than 1 year months Sklearn- Exprience - Less than 1 year months Scipy- Exprience - Less than 1 year months Keras- Exprience - Less than 1 year monthsCompany Details company - Matelabs description - ML Platform for business professionals, dummies and enthusiasts. 60/A Koramangala 5th block, Achievements/Tasks behind sukh sagar, Bengaluru, India                               Developed and deployed auto preprocessing steps of machine learning mainly missing value treatment, outlier detection, encoding, scaling, feature selection and dimensionality reduction. Deployed automated classification and regression model. linkedin.com/in/aditya-rathore- b4600b146                           Reasearch and deployed the time series forecasting model ARIMA, SARIMAX, Holt-winter and Prophet. Worked on meta-feature extracting problem. github.com/rathorology Implemented a state of the art research paper on outlier detection for mixed attributes. company - Matelabs description - "
            ],
            'augmented text': [
                "Phone: Not specified\nEmail: Not specified\nEducation: - May 2013 to May 2017: B.E, UIT-RGPV\nExperience: - Data Scientist, Matelabs (May 2017 - Present)\nSkills: - Python: Less than 1 year experience - Statsmodels: 12 months experience - AWS: Less than 1 year experience - Machine learning: Less than 1 year experience - Scikit-learn: Less than 1 year experience - Scipy: Less than 1 year experience - Keras: Less than 1 year experience\nProjects: - Developed and deployed auto preprocessing steps of machine learning at Matelabs - Deployed automated classification and regression model at Matelabs - Research and deployed time series forecasting models ARIMA, SARIMAX, Holt-winter and Prophet at Matelabs - Implemented a state of the art research paper on outlier detection for mixed attributes on github.com/rathorology\nTools: - Not specified Note: The company name \"Matelabs\" is mentioned multiple times, so it is assumed to be the same organization for all experience and project entries."
            ]
        })
        print("Using dummy data for demonstration. Training will be ineffective without a full dataset.")

    def format_to_instruction(example):
        instruction = (
            f"You are an expert Resume Information Extractor. Your task is to extract structured information "
            f"from the raw resume text provided below and format it into a detailed, structured representation. "
            f"The category of the resume is '{example['category']}'.\n\n"
            f"RAW RESUME TEXT:\n{example['resume'].strip()}\n\n"
            f"Extracted Structured Data:"
        )

        return {
            "text": f"[INST] {instruction} [/INST]\n{example['augmented text'].strip()}"
        }

    # Convert DataFrame to Hugging Face Dataset
    dataset = Dataset.from_pandas(df)
    
    # Apply the formatting function
    dataset = dataset.map(format_to_instruction, remove_columns=df.columns.tolist())
    
    print("Example formatted training sample (first 500 characters):")
    print("="*50)
    print(dataset[0]['text'][:500] + "...")
    print("="*50)
    
    return dataset

# --- 3. QLoRA Model Setup ---

def setup_model_and_tokenizer():
    # 4-bit quantization configuration
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",  # Normalized Float 4
        bnb_4bit_compute_dtype=torch.bfloat16, # Recommended for performance
        bnb_4bit_use_double_quant=True,
    )

    # Load the base model with QLoRA configuration
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=torch.bfloat16,
    )
    
    # Prepare model for k-bit training (important for QLoRA)
    model.config.use_cache = False
    model.config.pretraining_tp = 1
    model = prepare_model_for_kbit_training(model)

    # Load Tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    # Set padding token to EOS token for Mistral
    tokenizer.pad_token = tokenizer.eos_token 
    tokenizer.padding_side = "right" 

    # LoRA configuration
    # Target all linear layers in Mistral for better instruction-tuning performance
    peft_config = LoraConfig(
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        r=LORA_R,
        bias="none",
        task_type="CAUSAL_LM",
        # Targetting all linear layers for comprehensive instruction-tuning
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] 
    )
    
    return model, tokenizer, peft_config

# --- 4. Training Execution ---

def train_model():
    # 4.1 Load Data
    train_dataset = load_and_format_data(CSV_FILE_PATH)
    
    # 4.2 Setup Model, Tokenizer, and LoRA
    model, tokenizer, peft_config = setup_model_and_tokenizer()

    # 4.3 Setup Training Arguments
    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=NUM_TRAIN_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
        optim="paged_adamw_8bit", # Recommended optimizer for QLoRA
        save_steps=100,
        logging_steps=25,
        learning_rate=LEARNING_RATE,
        weight_decay=0.001,
        fp16=False,
        bf16=True if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8 else False,
        max_grad_norm=0.3,
        warmup_ratio=0.03,
        lr_scheduler_type="constant",
    )

    # 4.4 Initialize SFT Trainer
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        peft_config=peft_config,
        dataset_text_field="text", # The column we created with the formatted data
        tokenizer=tokenizer,
        max_seq_length=MAX_SEQ_LENGTH,
        packing=True, # Packs multiple samples into one sequence for efficiency
    )

    # 4.5 Start Training
    print("\nStarting QLoRA Fine-Tuning...")
    trainer.train()

    # 4.6 Save the fine-tuned adapter weights
    trainer.model.save_pretrained(os.path.join(OUTPUT_DIR, "final_checkpoint"))
    tokenizer.save_pretrained(os.path.join(OUTPUT_DIR, "final_checkpoint"))
    print(f"\nTraining complete. LoRA adapter saved to {os.path.join(OUTPUT_DIR, 'final_checkpoint')}")
    
    # Optional: Merge the adapter weights with the base model for easy deployment
    # This requires sufficient CPU RAM.
    # print("Merging adapter and saving full model...")
    # merged_model = trainer.model.merge_and_unload()
    # merged_model.save_pretrained(os.path.join(OUTPUT_DIR, "merged_model"))
    # tokenizer.save_pretrained(os.path.join(OUTPUT_DIR, "merged_model"))

if __name__ == "__main__":
    if torch.cuda.is_available():
        print(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
        train_model()
    else:
        print("CUDA not available. Please run this script on a GPU instance.")