from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset

# Definir el nombre del modelo
model_name = "bert-base-uncased"

# Cargar el tokenizador y el modelo preentrenado
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# Cargar el dataset de Hugging Face
dataset = load_dataset("GenaroCoronel/godot_docs")

print("Data set cargado godot_docs.json .....")
print(dataset)

# Verifica las columnas originales antes de cualquier procesamiento
print("Columnas originales del dataset:")
print(dataset['train'].column_names)

# Función para concatenar los campos relevantes
def preprocess_function(examples):
    texts = []
    for i in range(len(examples['name'])):
        text = f"Name: {examples['name'][i]}, Inherits: {examples['inherits'][i]}, "
        text += f"Brief Description: {examples['brief_description'][i]}, Description: {examples['description'][i]}"
        
        # Concatenar métodos
        methods = examples['methods'][i]
        methods_text = "Methods: "
        for method in methods:
            methods_text += f"{method['name']}({', '.join([param['name'] for param in method['params']])}), "
        text += methods_text
        
        # Concatenar miembros
        members = examples['members'][i]
        members_text = "Members: "
        for member in members:
            members_text += f"{member['name']} (Type: {member['type']}), "
        text += members_text
        
        texts.append(text)
    
    return {"text": texts}

# Preprocesar el dataset
processed_datasets = dataset.map(preprocess_function, batched=True)

# Verifica las columnas del dataset procesado
print("Columnas después de preprocesar:")
print(processed_datasets['train'].column_names)

# Tokenizar el dataset
def tokenize_function(examples):
    # Verifica si la columna 'text' está presente
    if 'text' not in examples:
        print("Columna 'text' no encontrada, saltando esta iteración.")
        return {}  # Si no hay 'text', no retorna nada para esta iteración
    
    # Tokenizar los textos
    tokenized_examples = tokenizer(examples['text'], padding="max_length", truncation=True)
    
    # Crear etiquetas basadas en algún criterio (ejemplo: 0 si 'Label' está en el nombre, 1 en otro caso)
    tokenized_examples['labels'] = [0 if 'Label' in text else 1 for text in examples['text']]
    
    return tokenized_examples

# Aplicar la tokenización al dataset
tokenized_datasets = processed_datasets.map(tokenize_function, batched=True)

# Filtrar filas que no contienen datos válidos
tokenized_datasets = tokenized_datasets.filter(lambda example: example is not None)

# Verifica que el dataset ha sido tokenizado correctamente
print("Dataset después de la tokenización:")
print(tokenized_datasets)

# Configuración de los argumentos de entrenamiento
training_args = TrainingArguments(
    output_dir='./results',          # Directorio donde guardar los resultados
    num_train_epochs=3,              # Número de épocas
    per_device_train_batch_size=8,   # Tamaño del batch para entrenamiento
    per_device_eval_batch_size=16,   # Tamaño del batch para evaluación
    logging_dir='./logs',            # Directorio para logs
)

# Inicializar el Trainer
trainer = Trainer(
    model=model,                     # El modelo que entrenaremos
    args=training_args,              # Argumentos de entrenamiento
    train_dataset=tokenized_datasets['train'],   # Dataset de entrenamiento
)

# Entrenar el modelo
trainer.train()

# Guardar el modelo entrenado
model.save_pretrained('./trained_model')
tokenizer.save_pretrained('./trained_model')

print("Entrenamiento completado y modelo guardado.")