import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer, GPT2Config
from torch.utils.data import Dataset, DataLoader

# Definir um pequeno conjunto de dados fictício
class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

# Exemplo de dados
training_data = [
"Usuário: O que é inteligência artificial?",
    "Modelo: Inteligência artificial refere-se à capacidade de uma máquina aprender e tomar decisões sem ser explicitamente programada.",
    "Usuário: Como as redes neurais funcionam?",
    "Modelo: As redes neurais são modelos inspirados no funcionamento do cérebro humano, compostos por camadas de neurônios interconectados.",
    "Usuário: olá, pode me ajudar?",
    "Modelo: Claro, em que posso ajudar?",
    "Usuário: Qual é o seu nome?",
    "Modelo: Meu nome é Theo"
    ]

# Criar uma instância do tokenizador GPT-2 e preparar os dados
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenized_data = tokenizer(training_data, return_tensors="pt", padding=True, truncation=True)

# Criar um DataLoader
dataset = CustomDataset(tokenized_data)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

# Configurar o modelo GPT-2 para treinamento
config = GPT2Config.from_pretrained("gpt2")
model = GPT2LMHeadModel(config)
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

# Treinar o modelo por algumas épocas (isso é um exemplo básico, ajuste conforme necessário)
num_epochs = 3
for epoch in range(num_epochs):
    for batch in dataloader:
        optimizer.zero_grad()
        outputs = model(**batch, labels=batch["input_ids"])
        loss = outputs.loss
        loss.backward()
        optimizer.step()

# Salvar o modelo ajustado
model.save_pretrained("modelo_gpt2_finetuned")

print("Treinamento concluído. Modelo salvo.")