#!/usr/bin/env python3
"""
Math Problem Solver Trainer
"""
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer

def train_math_model():
    print("🎓 Training Bengali Math Solver...")
    
    # Load data
    ds = load_dataset("hamim-87/Ashrafur_bangla_math", split="train[:10000]")
    
    # Initialize model
    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
    tokenizer.pad_token = tokenizer.eos_token
    
    # Prepare data
    def prepare_data(examples):
        texts = []
        for problem, solution in zip(examples['problem'], examples['solution']):
            text = f"প্রশ্ন: {problem}\n\nউত্তর: {solution}\n\n"
            texts.append(text)
        return tokenizer(texts, truncation=True, padding=True, max_length=512)
    
    tokenized_ds = ds.map(prepare_data, batched=True)
    
    # Training
    training_args = TrainingArguments(
        output_dir="./bangla_math_model",
        num_train_epochs=2,
        per_device_train_batch_size=4,
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_ds,
    )
    
    trainer.train()
    trainer.save_model()
    print("✅ Math model trained!")

if __name__ == "__main__":
    train_math_model()