Andro0s commited on
Commit
a3d0354
·
verified ·
1 Parent(s): 5d78646

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +285 -56
app.py CHANGED
@@ -1,63 +1,292 @@
1
- # Importamos las clases necesarias de la librería Hugging Face Transformers.
2
- # Usamos GPT2LMHeadModel porque incluye el "Language Model Head"
3
- # necesario para tareas de generación de texto.
4
- from transformers import GPT2LMHeadModel, GPT2Tokenizer
5
- import torch
6
- import sys
7
-
8
- # --- CONFIGURACIÓN ESTÁNDAR ---
9
- # MODEL_ID: Usamos 'gpt2', el identificador oficial y compatible.
10
- # ESTA ES LA CORRECCIÓN CLAVE para evitar el error de compatibilidad.
11
- MODEL_ID = 'gpt2'
12
- PROMPT = "La arquitectura de la red neuronal Transformer se basa en un mecanismo de"
13
- MAX_LENGTH = 70
14
- TEMPERATURE = 0.8
15
- DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
16
-
17
- print(f"Iniciando carga del modelo GPT-2: {MODEL_ID} en dispositivo: {DEVICE}")
18
-
19
- def main():
20
- """Función principal para cargar el modelo y generar texto."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
21
  try:
22
- # 1. Cargar el Tokenizer
23
- tokenizer = GPT2Tokenizer.from_pretrained(MODEL_ID)
24
-
25
- # 2. Cargar el Modelo con el 'Language Modeling Head'
26
- model = GPT2LMHeadModel.from_pretrained(MODEL_ID).to(DEVICE)
27
-
28
- # Configuramos el token de padding (necesario para la generación)
29
- if tokenizer.pad_token is None:
30
- tokenizer.pad_token = tokenizer.eos_token
31
- model.config.pad_token_id = tokenizer.eos_token_id
32
-
33
- # 3. Codificar el texto de inicio (prompt)
34
- input_ids = tokenizer.encode(PROMPT, return_tensors='pt').to(DEVICE)
35
-
36
- print("\n--- Generando Texto ---")
37
-
38
- # 4. Generar la secuencia de texto
39
- output = model.generate(
40
- input_ids,
41
- max_length=MAX_LENGTH,
42
- num_return_sequences=1,
43
- do_sample=True, # Permite sampling (generación más creativa)
44
- top_k=50, # Limita las opciones a las 50 más probables
45
- temperature=TEMPERATURE, # Controla la creatividad (más alto = más aleatorio)
46
- pad_token_id=tokenizer.eos_token_id
47
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48
 
49
- # 5. Decodificar y mostrar el resultado
50
- generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
 
 
 
 
 
 
 
51
 
52
- print(f"\nPrompt Original: {PROMPT}")
53
- print("-" * (len(PROMPT) + 15))
54
- print(f"Texto Generado Completo: {generated_text}")
55
- print("\n¡Generación completada con éxito!")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56
 
57
  except Exception as e:
58
- print(f"\n[ERROR CRÍTICO] Fallo al cargar el modelo o generar texto.", file=sys.stderr)
59
- print(f"Asegúrate de tener instalada la librería 'transformers' y 'torch' (pip install transformers torch).", file=sys.stderr)
60
- print(f"Detalle del error: {e}", file=sys.stderr)
 
 
 
 
 
 
 
 
 
 
 
 
61
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
62
  if __name__ == "__main__":
63
- main()
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ import threading
4
+ import time
5
+ from huggingface_hub import login
6
+ from datasets import load_dataset
7
+ from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, pipeline
8
+ from peft import get_peft_model, LoraConfig, TaskType, PeftModel
9
+ import json
10
+
11
+ # --- CONFIGURACIÓN DEL MODELO Y ENTRENAMIENTO ---
12
+ BASE_MODEL = "bigcode/santacoder" # Modelo base de programación
13
+ LORA_PATH = "./lora_output" # Ruta donde se guarda el modelo adaptado
14
+ DATASET_FILE = "codesearchnet_lora_dataset.json"
15
+ MAX_TOKEN_LENGTH = 256
16
+ NUM_SAMPLES_TO_PROCESS = 1000
17
+ DEFAULT_EPOCHS = 10
18
+
19
+ # Configuración del ciclo AUTÓNOMO (Inicia reentrenamiento cada 5 interacciones)
20
+ GENERATION_LIMIT_TO_TRAIN = 5
21
+ AUTONOMOUS_EPOCHS = 3
22
+
23
+ # --- ESTADO GLOBAL Y THREADING ---
24
+ tokenizer = None
25
+ lora_model = None
26
+ tokenized_dataset = None
27
+ lora_generator = None
28
+
29
+ # Variables de estado
30
+ version_number = 1.0
31
+ is_trained = os.path.exists(LORA_PATH)
32
+ generations_since_last_train = 0
33
+ training_status_message = "Esperando la inicialización V1.0..."
34
+
35
+ # Lock para proteger las variables compartidas entre hilos (CRÍTICO para estabilidad)
36
+ global_lock = threading.Lock()
37
+
38
+ # --- LÓGICA DE PREPARACIÓN Y SETUP ---
39
+
40
+ def prepare_codesearchnet():
41
+ """Descarga y prepara el dataset inicial si no existe."""
42
+ if os.path.exists(DATASET_FILE):
43
+ return
44
  try:
45
+ raw_csn = load_dataset('Nan-Do/code-search-net-python', split=f'train[:{NUM_SAMPLES_TO_PROCESS}]')
46
+
47
+ def format_for_lora(example):
48
+ # Formato que entrena a la IA a enlazar descripción (español) con código (inglés)
49
+ prompt_text = (
50
+ f"# Descripción: {example['docstring_summary']}\n"
51
+ f"# Completa la siguiente función:\n"
52
+ f"def {example['func_name']}("
53
+ )
54
+ completion_text = example['code']
55
+ return {"prompt": prompt_text, "completion": completion_text}
56
+
57
+ lora_dataset = raw_csn.map(format_for_lora, batched=False, remove_columns=raw_csn["train"].column_names)
58
+ lora_dataset.to_json(DATASET_FILE)
59
+ except Exception as e:
60
+ print(f"Error al cargar dataset. Usando datos mínimos. Error: {e}")
61
+ minimal_dataset = [{"prompt": "# Error de carga. Intenta de nuevo.", "completion": "pass\n"}] * 10
62
+ with open(DATASET_FILE, 'w') as f:
63
+ json.dump(minimal_dataset, f)
64
+
65
+
66
+ def setup_resources():
67
+ """Configura el tokenizer, el modelo base y el adaptador LoRA."""
68
+ global tokenizer, lora_model, tokenized_dataset
69
+
70
+ prepare_codesearchnet()
71
+
72
+ hf_token = os.environ.get("HF_TOKEN")
73
+ if hf_token:
74
+ login(token=hf_token)
75
+
76
+ tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
77
+ base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map="auto")
78
+
79
+ if tokenizer.pad_token is None:
80
+ tokenizer.pad_token = tokenizer.eos_token
81
+
82
+ peft_config = LoraConfig(
83
+ task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=32, lora_dropout=0.1, target_modules=["c_proj", "c_attn"],
84
+ )
85
+ lora_model = get_peft_model(base_model, peft_config)
86
+
87
+ try:
88
+ raw_dataset = load_dataset("json", data_files=DATASET_FILE)
89
+
90
+ def tokenize_function(examples):
91
+ return tokenizer(
92
+ examples["prompt"] + examples["completion"],
93
+ truncation=True,
94
+ padding="max_length",
95
+ max_length=MAX_TOKEN_LENGTH
96
+ )
97
+
98
+ tokenized_dataset = raw_dataset.map(tokenize_function, batched=True, remove_columns=raw_dataset["train"].column_names if "train" in raw_dataset else [],)
99
+ except Exception:
100
+ tokenized_dataset = None
101
+
102
+ # --- FUNCIÓN DE ENTRENAMIENTO (EJECUTADA EN HILO SEPARADO) ---
103
+
104
+ def autonomous_train_lora(epochs, batch_size, learning_rate):
105
+ """Ejecuta el entrenamiento en un hilo separado para la autonomía."""
106
+ global lora_model, tokenized_dataset, lora_generator, version_number, is_trained, training_status_message
107
+
108
+ try:
109
+ with global_lock:
110
+ if tokenized_dataset is None or "train" not in tokenized_dataset:
111
+ training_status_message = "ERROR: No se puede entrenar. Dataset no disponible."
112
+ return
113
+
114
+ # 1. ACTUALIZAR VERSIÓN (Pre-incremento)
115
+ if is_trained:
116
+ version_number += 0.1
117
+ else:
118
+ version_number = 1.0
119
+
120
+ # 2. CONFIGURACIÓN E INICIO DEL ENTRENAMIENTO
121
+ training_status_message = f"🧠 ENTRENANDO V{version_number:.1f} (Epochs: {epochs})...."
122
+ print(f"\n[AUTÓNOMO] {training_status_message}")
123
+
124
+ data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
125
+ training_args = TrainingArguments(
126
+ output_dir=LORA_PATH,
127
+ per_device_train_batch_size=int(batch_size),
128
+ num_train_epochs=float(epochs),
129
+ learning_rate=float(learning_rate),
130
+ save_total_limit=1,
131
+ logging_steps=10,
132
+ push_to_hub=False,
133
+ disable_tqdm=True,
134
+ report_to="none"
135
+ )
136
+
137
+ trainer = Trainer(model=lora_model, args=training_args, train_dataset=tokenized_dataset["train"], data_collator=data_collator)
138
+
139
+ trainer.train()
140
+ lora_model.save_pretrained(LORA_PATH)
141
+ tokenizer.save_pretrained(LORA_PATH)
142
+
143
+ # 3. Marcar como entrenado
144
+ is_trained = True
145
+ training_status_message = f"✅ ENTRENAMIENTO V{version_number:.1f} COMPLETADO. Modelo listo para Hot Swap."
146
+ print(f"[AUTÓNOMO] {training_status_message}")
147
+
148
+ except Exception as e:
149
+ training_status_message = f"ERROR CRÍTICO durante el entrenamiento autónomo: {e}"
150
+ print(f"[AUTÓNOMO] {training_status_message}")
151
+
152
+ # --- FUNCIÓN DE GENERACIÓN (CORREGIDA PARA RETORNAR 2 VALORES) ---
153
+
154
+ def generate_text(prompt_text):
155
+ """Genera código y dispara el ciclo de reentrenamiento autónomo si es necesario."""
156
+ global lora_generator, generations_since_last_train, is_trained, version_number, training_status_message
157
+
158
+ if not is_trained:
159
+ # Si el entrenamiento V1.0 no ha terminado, retorna el mensaje de error y el estado actual
160
+ return "ERROR: El modelo LoRA no ha sido entrenado. Por favor, espere mientras la IA se inicializa con el entrenamiento V1.0.", update_status()
161
+
162
+
163
+ # 1. HOT SWAP (Verifica si el modelo necesita recargarse con la nueva versión)
164
+ if lora_generator is None:
165
+ with global_lock:
166
+ try:
167
+ # Recarga el modelo solo si está vacío
168
+ base_model_gen = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map="auto")
169
+ model_with_lora = PeftModel.from_pretrained(base_model_gen, LORA_PATH)
170
+ final_model = model_with_lora.merge_and_unload()
171
+ final_model.eval()
172
+ lora_generator = pipeline("text-generation", model=final_model, tokenizer=tokenizer)
173
+ print(f"[HOT SWAP] 🔄 Modelo de inferencia V{version_number:.1f} recargado y listo.")
174
+ except Exception as e:
175
+ # Si la recarga falla, retorna un error
176
+ return f"Error al cargar el modelo V{version_number:.1f} para inferencia: {e}", update_status()
177
+
178
 
179
+ # 2. Generación de texto (Lógica de inferencia)
180
+ try:
181
+ # Prepara el prompt para guiar la generación del código
182
+ prompt_with_indent = prompt_text.strip() + "\n "
183
+ output = lora_generator(prompt_with_indent, max_new_tokens=150, temperature=0.7, top_p=0.9, clean_up_tokenization_spaces=True)
184
+ full_output = output[0]["generated_text"]
185
+
186
+ # Extrae solo la parte de la compleción (el código generado)
187
+ start_index = full_output.find(prompt_with_indent)
188
+ completion = full_output[start_index + len(prompt_with_indent):] if start_index != -1 else full_output
189
 
190
+ # 3. Aumentar contador de autonomía
191
+ with global_lock:
192
+ generations_since_last_train += 1
193
+ current_count = generations_since_last_train
194
+ current_version = version_number
195
+
196
+ # 4. Verificar si se requiere reentrenamiento (y dispararlo en un nuevo hilo)
197
+ notification = ""
198
+ if current_count >= GENERATION_LIMIT_TO_TRAIN:
199
+ # Verifica que no haya otro hilo de entrenamiento ya corriendo
200
+ if not any(isinstance(t, threading.Thread) and t.name == 'AutonomousTrainer' for t in threading.enumerate()):
201
+ print(f"[AUTONOMÍA] Generación #{current_count} alcanzada. Disparando reentrenamiento autónomo en segundo plano...")
202
+
203
+ # Reiniciar el contador de generaciones y forzar Hot Swap en la próxima interacción
204
+ with global_lock:
205
+ generations_since_last_train = 0
206
+ lora_generator = None
207
+
208
+ trainer_thread = threading.Thread(
209
+ target=autonomous_train_lora,
210
+ args=(AUTONOMOUS_EPOCHS, 2, 5e-5),
211
+ name='AutonomousTrainer'
212
+ )
213
+ trainer_thread.daemon = True
214
+ trainer_thread.start()
215
+
216
+ notification = f"\n\n--- [AUTONOMÍA] La IA ha iniciado el reentrenamiento V{current_version+0.1:.1f} para mejorar la traducción de tu diálogo. La próxima generación cargará la nueva versión. ---"
217
+
218
+ # CORRECCIÓN CLAVE: Retorna el código Y el estado actualizado
219
+ return completion + notification, update_status()
220
 
221
  except Exception as e:
222
+ # Si falla la generación, retorna el mensaje de error y el estado actual
223
+ return f"Error generando texto: {e}", update_status()
224
+
225
+ # --- FUNCIÓN PARA INICIALIZACIÓN Y ENTRENAMIENTO V1.0 (Obligatorio) ---
226
+
227
+ def initialize_and_train_v1():
228
+ """Ejecuta el entrenamiento inicial V1.0 de forma autónoma al iniciar."""
229
+ if not is_trained:
230
+ autonomous_train_lora(epochs=DEFAULT_EPOCHS, batch_size=2, learning_rate=5e-5)
231
+ else:
232
+ global training_status_message
233
+ training_status_message = f"✅ Modelo V{version_number:.1f} ya entrenado. Listo."
234
+ print(f"[INICIALIZACIÓN] {training_status_message}")
235
+
236
+ # --- FUNCIÓN PARA ACTUALIZAR EL ESTADO EN LA UI ---
237
 
238
+ def update_status():
239
+ """Actualiza la versión y el estado del entrenamiento en la interfaz de Gradio."""
240
+ global training_status_message, version_number
241
+ # Retorna un texto en Markdown que se actualiza constantemente
242
+ return f"**Versión de Comprensión:** V{version_number:.1f} | **Estado del Entrenador:** {training_status_message}"
243
+
244
+
245
+ # --- INTERFAZ GRADIO ---
246
+ with gr.Blocks(title="AmorCoderAI - Aprendizaje Continuo") as demo:
247
+ gr.Markdown("# 💙 AmorCoderAI - Asistente de Código con Aprendizaje Continuo")
248
+
249
+ # Muestra la versión y el estado.
250
+ version_and_status = gr.Markdown(
251
+ f"**Versión de Comprensión:** V{version_number:.1f} | **Estado del Entrenador:** {training_status_message}",
252
+ elem_id="status_display"
253
+ )
254
+
255
+ gr.Markdown(f"**Modo Autónomo:** La IA se reentrena automáticamente cada **{GENERATION_LIMIT_TO_TRAIN}** códigos generados. Esto mejora su capacidad para traducir tu español conversacional a código.")
256
+
257
+ with gr.Tab("✨ Generación de Código"):
258
+ gr.Markdown("## Escribe tu idea en palabras (¡Usa español fluido!)")
259
+
260
+ gr.Markdown("Recomendación inicial: Usa el siguiente formato para obtener el mejor código mientras la IA aprende tu idioma:")
261
+
262
+ prompt = gr.Textbox(
263
+ label="Instrucción de Programación:",
264
+ lines=4,
265
+ placeholder="# Descripción: Quiero que me hagas un código similar a Google Gemini.\n# Completa la siguiente función:\ndef generar_contenido(prompt, modelo):"
266
+ )
267
+ generate_button = gr.Button("💬 Generar código y disparar Aprendizaje")
268
+ output_box = gr.Textbox(label="Código generado", lines=10)
269
+
270
+ # Conexión del botón con la función principal
271
+ # IMPORTANTE: Ahora generate_text retorna DOS valores para coincidir con [output_box, version_and_status]
272
+ generate_button.click(
273
+ generate_text,
274
+ inputs=prompt,
275
+ outputs=[output_box, version_and_status],
276
+ )
277
+
278
+ # El estado se actualiza solo al cargar la página.
279
+ demo.load(update_status, None, version_and_status)
280
+
281
+
282
+ # --- INICIO DE LA APLICACIÓN ---
283
  if __name__ == "__main__":
284
+ setup_resources()
285
+
286
+ # Lanza el entrenamiento V1.0 inicial en un hilo para que no congele la UI
287
+ initialization_thread = threading.Thread(target=initialize_and_train_v1, name='InitializationTrainer')
288
+ initialization_thread.daemon = True
289
+ initialization_thread.start()
290
+
291
+ print(f"\n💻 LANZANDO INTERFAZ GRADIO (El entrenamiento V1.0 se ejecuta en segundo plano)")
292
+ demo.launch()