transformer.py

# Importar las librerías necesarias
import transformers
import datasets
from datasets import load_dataset
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments

# 1. Cargar los datasets
# Estos datasets se utilizarán para el entrenamiento
ds_chatgpt_prompts = load_dataset("fka/awesome-chatgpt-prompts")
ds_open_thoughts = load_dataset("open-thoughts/OpenThoughts-114k")
ds_bangla_math = load_dataset("hamim-87/Ashrafur_bangla_math")
ds_deepseek_prover = load_dataset("deepseek-ai/DeepSeek-Prover-V1")
ds = load_dataset("facebook/natural_reasoning")


# Puedes agregar más datasets de la siguiente manera
# ds_nuevo = load_dataset("nombre_del_dataset")

# 2. Usar pipeline para tareas de generación de texto con un modelo preentrenado
# Esto es útil para realizar tareas como texto generativo, traducción o clasificación

# Definir el modelo y la tarea para el pipeline
pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)

# Mensajes de ejemplo para generar texto
messages = [
    {"role": "user", "content": "Who are you?"}
]

# Generar texto con el pipeline
generated_response = pipe(messages)
print("Generación de texto:", generated_response)

# 3. Cargar el modelo de causal LM (Lenguaje de Modelos Causales)
# Usamos el modelo "DeepSeek-R1" como ejemplo
model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)

# Tokenización del texto de entrada
input_text = "What is the capital of France?"
inputs = tokenizer(input_text, return_tensors="pt")

# Generar texto con el modelo cargado
generated_outputs = model.generate(inputs['input_ids'], max_length=50)
generated_text = tokenizer.decode(generated_outputs[0], skip_special_tokens=True)
print("Generación con el modelo cargado:", generated_text)

# 4. Definir parámetros de entrenamiento con TrainingArguments
training_args = TrainingArguments(
    output_dir='./results',             # Directorio de salida para los resultados
    evaluation_strategy="epoch",        # Evaluar al final de cada época
    learning_rate=2e-5,                 # Tasa de aprendizaje
    per_device_train_batch_size=8,      # Tamaño del lote para entrenamiento
    per_device_eval_batch_size=16,      # Tamaño del lote para evaluación
    num_train_epochs=3,                 # Número de épocas
    weight_decay=0.01,                  # Regularización
)

# 5. Preparar el Trainer con el dataset y el modelo
# Seleccionamos el dataset y lo preprocesamos
# Ejemplo: usar "ds_chatgpt_prompts" para entrenamiento (aquí puedes modificar con otros datasets)

train_dataset = ds_chatgpt_prompts['train']
eval_dataset = ds_chatgpt_prompts['test']

# 6. Configurar el Trainer
trainer = Trainer(
    model=model,                         # El modelo que entrenaremos
    args=training_args,                  # Los argumentos de entrenamiento
    train_dataset=train_dataset,         # Conjunto de datos de entrenamiento
    eval_dataset=eval_dataset,           # Conjunto de datos de evaluación
)

# Entrenar el modelo
trainer.train()

# 7. Guardar el modelo entrenado
trainer.save_model('./my_trained_model')

# 8. Evaluación del modelo (opcional)
results = trainer.evaluate()
print("Evaluación del modelo:", results)

# 9. Uso de otros modelos para tareas específicas
# Aquí puedes seguir agregando otros modelos para tareas como traducción, clasificación, etc.

# Ejemplo: Traducción con mbart (traducción inglés a árabe)
translation_pipe = pipeline("translation", model="akhooli/mbart-large-cc25-ar-en", trust_remote_code=True)
translation_input = "Hello, how are you?"
translated_output = translation_pipe(translation_input)
print("Traducción:", translated_output)

# Ejemplo de usar otro modelo de texto de recuperación (por ejemplo, GTE Multilingual)
gte_pipe = pipeline("text-classification", model="Alibaba-NLP/gte-multilingual-reranker-base", trust_remote_code=True)
gte_input = "What are the benefits of AI?"
gte_output = gte_pipe(gte_input)
print("Clasificación GTE:", gte_output)

# Puedes seguir agregando más modelos y pipelines como estos