from transformers import GPT2LMHeadModel, TrainingArguments, Trainer, DataCollatorForLanguageModeling, GPT2Tokenizer
import gradio as gr

# Create a list of 30 emoji math problems with their solutions.
# Format: "Q: [emoji math equation]\nA: [solution]"
data = [
    "Q: 🍎 + 🍎 + 🍎 = 12\nA: 4",
    "Q: 🎲 + 🎲 = 12\nA: 6",
    "Q: 🚗 + 🚗 + 🚗 + 🚗 = 20\nA: 5",
    "Q: 🍌 + 🍌 + 🍌 + 🍌 + 🍌 = 15\nA: 3",
    "Q: 🍓 + 🍓 = 8\nA: 4",
    "Q: 🍕 + 🍕 + 🍕 = 18\nA: 6",
    "Q: 🍩 + 🍩 + 🍩 + 🍩 = 20\nA: 5",
    "Q: 🌟 + 🌟 + 🌟 = 9\nA: 3",
    "Q: 🎈 + 🎈 = 14\nA: 7",
    "Q: 🎂 + 🎂 + 🎂 = 15\nA: 5",
    "Q: 🍪 + 🍪 + 🍪 + 🍪 = 16\nA: 4",
    "Q: 🍭 + 🍭 + 🍭 = 15\nA: 5",
    "Q: 🧁 + 🧁 = 10\nA: 5",
    "Q: 🥑 + 🥑 + 🥑 = 12\nA: 4",
    "Q: 🍇 + 🍇 = 10\nA: 5",
    "Q: 🍒 + 🍒 + 🍒 = 15\nA: 5",
    "Q: 🍍 + 🍍 = 14\nA: 7",
    "Q: 🍉 + 🍉 + 🍉 + 🍉 = 20\nA: 5",
    "Q: 🥭 + 🥭 = 16\nA: 8",
    "Q: 🍈 + 🍈 + 🍈 = 9\nA: 3",
    "Q: 🍑 + 🍑 + 🍑 + 🍑 = 20\nA: 5",
    "Q: 🍏 + 🍏 = 10\nA: 5",
    "Q: 🍋 + 🍋 + 🍋 = 12\nA: 4",
    "Q: 🍊 + 🍊 = 10\nA: 5",
    "Q: 🥝 + 🥝 + 🥝 = 15\nA: 5",
    "Q: 🍐 + 🍐 = 8\nA: 4",
    "Q: 🍆 + 🍆 + 🍆 + 🍆 = 16\nA: 4",
    "Q: 🥕 + 🥕 = 10\nA: 5",
    "Q: 🌽 + 🌽 + 🌽 = 9\nA: 3",
    "Q: 🥔 + 🥔 + 🥔 + 🥔 = 20\nA: 5"
]

# For training with Hugging Face's datasets, we create a dictionary.
from datasets import Dataset
dataset = Dataset.from_dict({"text": data})

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize_function(example):
    return tokenizer(example["text"], truncation=True, max_length=128, padding="max_length")

tokenized_dataset = dataset.map(tokenize_function, batched=True)
tokenized_dataset.set_format(type="torch", columns=["input_ids", "attention_mask"])

# Load GPT-2 model
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.config.pad_token_id = tokenizer.eos_token_id

# Create a data collator for language modeling that will handle padding dynamically
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

# Define training arguments
training_args = TrainingArguments(
    output_dir="./emoji_math_model",
    overwrite_output_dir=True,
    num_train_epochs=8,
    per_device_train_batch_size=4,
    save_steps=50,
    save_total_limit=2,
    logging_steps=10,
    learning_rate=1e-5
)

# Initialize the Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

# Start training
trainer.train()

import logging
logging.getLogger("transformers").setLevel(logging.ERROR)  # Suppress transformer warnings

import re

def generate_single_answer(prompt, max_new_tokens=10):
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to(model.device)
    attention_mask = inputs["attention_mask"].to(model.device)

    output = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_new_tokens=max_new_tokens,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=False,
        repetition_penalty=2.0
    )

    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

    if prompt in generated_text:
        generated_text = generated_text.split(prompt, 1)[1]

    match = re.search(r'\b(\d+)\b', generated_text)
    if match:
        answer = match.group(1)
    else:
        answer = generated_text.strip()

    return answer

# Gradio UI
def gradio_interface(prompt):
    return generate_single_answer(prompt)

iface = gr.Interface(fn=gradio_interface, inputs="text", outputs="text", title="Emoji Math Solver", description="Enter an emoji math problem to get the answer.")
iface.launch()