import spaces
import gradio as gr
from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM, TrainerCallback
from datasets import load_dataset
import traceback
from huggingface_hub import login
from peft import get_peft_model, LoraConfig

class LoggingCallback(TrainerCallback):
    def on_step_end(self, args, state, control, kwargs):
        # Log the learning rate
        current_lr = state.optimizer.param_groups[0]['lr']
        print(f"Current Learning Rate: {current_lr}")

    def on_epoch_end(self, args, state, control, kwargs):
        # Log the error rate (assuming you have a metric to calculate it)
        # Here we assume you have a way to get the validation loss
        if state.best_metric is not None:
            error_rate = 1 - state.best_metric  # Assuming best_metric is accuracy
            print(f"Current Error Rate: {error_rate:.4f}")

@spaces.GPU
def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
    try:
        login(api_key.strip())
        lora_config = LoraConfig(
            r=16,  # Rank of the low-rank adaptation
            lora_alpha=32,  # Scaling factor
            lora_dropout=0.1,  # Dropout for LoRA layers
            bias="none"  # Bias handling
        )        
        # Load the dataset
        dataset = load_dataset(dataset_name.strip())
    
        # Load the model and tokenizer
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
        model = get_peft_model(model, lora_config)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
    
        # Tokenize the dataset
        def tokenize_function(examples):
            max_length = 256
            return tokenizer(examples['text'], max_length=max_length, truncation=True)
    
        tokenized_datasets = dataset.map(tokenize_function, batched=True)
    
        # Set training arguments
        training_args = TrainingArguments(
            output_dir='./results',
            eval_strategy="epoch",
            save_strategy='epoch', 
            learning_rate=lr*0.00001,
            per_device_train_batch_size=int(batch_size),
            per_device_eval_batch_size=int(batch_size), 
            num_train_epochs=int(num_epochs),
            weight_decay=0.01,
            gradient_accumulation_steps=grad*0.1,
            load_best_model_at_end=True,
            metric_for_best_model="accuracy",
            greater_is_better=True,
            logging_dir='./logs',
            logging_steps=10,
            #push_to_hub=True,
            hub_model_id=hub_id.strip(),
            fp16=True,
            lr_scheduler_type='cosine',
        )
        
        # Create Trainer
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized_datasets['train'],
            eval_dataset=tokenized_datasets['test'],
            test_dataset=tokenized_datasets['validation'],
            callbacks=[LoggingCallback()], 
        )
    
        # Fine-tune the model
        trainer.train()
        trainer.push_to_hub(commit_message="Training complete!")
    except Exception as e:
        return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"
    return 'DONE!'#model
'''
# Define Gradio interface
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = model(inputs)
    predictions = outputs.logits.argmax(dim=-1)
    return "Positive" if predictions.item() == 1 else "Negative"
'''
# Create Gradio interface
try:
    
    iface = gr.Interface(
        fn=fine_tune_model,
        inputs=[
            gr.Textbox(label="Model Name (e.g., 'google/t5-efficient-tiny-nh8')"),
            gr.Textbox(label="Dataset Name (e.g., 'imdb')"),
            gr.Textbox(label="HF hub to push to after training"),
            gr.Textbox(label="HF API token"),
            gr.Slider(minimum=1, maximum=10, value=3, label="Number of Epochs"),
            gr.Slider(minimum=1, maximum=16, value=4, label="Batch Size"),
            gr.Slider(minimum=1, maximum=1000, value=50, label="Learning Rate (e-5)"),
            gr.Slider(minimum=1, maximum=100, value=1, label="Gradient accumulation (e-1)"), 
        ],
        outputs="text",
        title="Fine-Tune Hugging Face Model",
        description="This interface allows you to fine-tune a Hugging Face model on a specified dataset."
    )
    # Launch the interface
    iface.launch()    
except Exception as e:
    print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")