import gradio as gr
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments

# Dataset loading (replace with your desired dataset)
dataset = load_dataset("meta-llama/Meta-Llama-3.1-8B-Instruct-evals", "Meta-Llama-3.1-8B-Instruct-evals__arc_challenge__details")

# Model and tokenizer (replace with desired model)
model_name = "mradermacher/llama-3-8b-gpt-4o-GGUF"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Training function (optional)
def train_model(epochs=3):
    training_args = TrainingArguments(
        output_dir="output",  # Adjust output directory
        per_device_train_batch_size=8,  # Adjust batch size
        num_train_epochs=epochs,
        evaluation_strategy="epoch",  # Adjust evaluation strategy
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
    )

    trainer.train()
    print("Model training complete!")

# Text generation function
def generate_text(prompt):
    try:
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
        output = model.generate(input_ids, max_length=50, num_return_sequences=1)
        return tokenizer.decode(output[0], skip_special_tokens=True)
    except Exception as e:
        return f"Error generating text: {e}"

# Gradio interface for text generation
interface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="Text Generation with Trained Model",
    description="Enter a prompt and get creative text generated by the model.",
)

# Train the model before launching the interface (optional)
train_model()  # Uncomment to train before launching

# Launch the Gradio interface
interface.launch()