app.py

import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer, BartForConditionalGeneration, BartTokenizer
import torch

# Configurar el dispositivo (CPU)
device = torch.device("cpu")

# Cargar el modelo y tokenizer para autocompletar código
print("Cargando modelo code-autocomplete-gpt2-base...")
model_name = "shibing624/code-autocomplete-gpt2-base"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

# Mover modelo a CPU y ponerlo en modo evaluación
model.to(device)
model.eval()

# Configurar pad_token si no existe
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Cargar BART para simplificación de texto
print("Cargando modelo BART para simplificación...")
bart_model_name = "facebook/bart-base"
bart_tokenizer = BartTokenizer.from_pretrained(bart_model_name)
bart_model = BartForConditionalGeneration.from_pretrained(bart_model_name)

# Mover BART a CPU y ponerlo en modo evaluación
bart_model.to(device)
bart_model.eval()

def autocomplete_text(input_text, max_tokens=20):
    """
    Autocompleta el texto/código de entrada usando code-autocomplete-gpt2-base
    
    Args:
        input_text (str): Texto/código inicial a completar
        max_tokens (int): Número máximo de tokens a generar
    
    Returns:
        str: Solo la parte nueva generada (sin el input original)
    """
    if not input_text.strip():
        return "Por favor, ingresa algún texto para completar."
    
    try:
        # Tokenizar el texto de entrada
        inputs = tokenizer.encode(input_text, return_tensors="pt", padding=True)
        inputs = inputs.to(device)
        
        # Generar texto
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=max_tokens,
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                attention_mask=torch.ones_like(inputs)
            )
        
        # Decodificar el resultado completo
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extraer solo la parte nueva (sin el input original)
        new_text = generated_text[len(input_text):].strip()
        
        if not new_text:
            return "No se pudo generar texto adicional."
        
        return new_text
        
    except Exception as e:
        return f"Error al generar texto: {str(e)}"

def simplify_text(input_text, max_length=150):
    """
    Simplifica texto complejo usando BART
    
    Args:
        input_text (str): Texto complejo a simplificar
        max_length (int): Longitud máxima del texto simplificado
    
    Returns:
        str: Texto simplificado con palabras más sencillas
    """
    if not input_text.strip():
        return "Por favor, ingresa algún texto para simplificar."
    
    try:
        # Crear un prompt para guiar la simplificación
        prompt = f"Simplify this text using easier words: {input_text}"
        
        # Tokenizar el texto
        inputs = bart_tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
        inputs = inputs.to(device)
        
        # Generar texto simplificado
        with torch.no_grad():
            outputs = bart_model.generate(
                inputs,
                max_length=max_length,
                min_length=20,
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                early_stopping=True,
                no_repeat_ngram_size=2,
                pad_token_id=bart_tokenizer.pad_token_id,
                eos_token_id=bart_tokenizer.eos_token_id
            )
        
        # Decodificar el resultado
        simplified_text = bart_tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Limpiar el texto (remover el prompt si aparece)
        if simplified_text.startswith("Simplify this text using easier words:"):
            simplified_text = simplified_text.replace("Simplify this text using easier words:", "").strip()
        
        if not simplified_text:
            return "No se pudo simplificar el texto."
        
        return simplified_text
        
    except Exception as e:
        return f"Error al simplificar texto: {str(e)}"

def create_autocomplete_interface():
    """
    Crea la interfaz con autocompletar y simplificación dentro de gr.Blocks()
    """
    
    with gr.Blocks(title="Asistente de Texto y Código") as demo:
        
        gr.Markdown("# 🤖 Asistente de Texto y Código")
        gr.Markdown("Herramientas para autocompletar código y simplificar textos complejos.")
        
        with gr.Tab("Autocompletar"):
            with gr.Row():
                with gr.Column():
                    input_textbox = gr.Textbox(
                        label="Código a completar",
                        placeholder="def fibonacci(n):",
                        lines=5,
                        max_lines=10
                    )
                    
                    generate_btn = gr.Button("Completar Código", variant="primary")
                    
                with gr.Column():
                    output_textbox = gr.Textbox(
                        label="Código generado",
                        placeholder="Aquí aparecerá la continuación del código...",
                        lines=5,
                        max_lines=10,
                        interactive=False
                    )
            
            # Conectar el botón con la función
            generate_btn.click(
                fn=autocomplete_text,
                inputs=[input_textbox],
                outputs=[output_textbox]
            )
            
            # También permitir Enter para generar
            input_textbox.submit(
                fn=autocomplete_text,
                inputs=[input_textbox],
                outputs=[output_textbox]
            )
            
        # Nueva pestaña para simplificar texto
        with gr.Tab("Simplificar Texto"):
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="Texto complejo a simplificar",
                        placeholder="Ingresa aquí el texto difícil de entender...",
                        lines=6,
                        max_lines=12
                    )
                    
                    simplify_btn = gr.Button("Simplificar Texto", variant="secondary")
                    
                with gr.Column():
                    simplified_output = gr.Textbox(
                        label="Texto simplificado",
                        placeholder="Aquí aparecerá el texto más fácil de entender...",
                        lines=6,
                        max_lines=12,
                        interactive=False
                    )
            
            # Conectar el botón de simplificar
            simplify_btn.click(
                fn=simplify_text,
                inputs=[text_input],
                outputs=[simplified_output]
            )
            
            # También permitir Enter para simplificar
            text_input.submit(
                fn=simplify_text,
                inputs=[text_input],
                outputs=[simplified_output]
            )
            
        # Pestaña adicional con ejemplos
        with gr.Tab("Ejemplos"):
            gr.Markdown("""
            ### Ejemplos de Autocompletado de Código:
            
            **Entrada:** "def fibonacci(n):"
            **Salida:** "\\n    if n <= 1:\\n        return n\\n    return fibonacci(n-1) + fibonacci(n-2)"
            
            **Entrada:** "for i in range("
            **Salida:** "10):\\n    print(i)"
            
            **Entrada:** "import pandas as pd\\ndf = pd.read_csv("
            **Salida:** "'data.csv')\\nprint(df.head())"
            
            ---
            
            ### Ejemplos de Simplificación de Texto:
            
            **Texto complejo:** "La implementación de algoritmos de machine learning requiere una comprensión profunda de estructuras de datos y técnicas de optimización."
            **Texto simple:** "Para usar inteligencia artificial necesitas entender bien cómo organizar datos y mejorar programas."
            
            **Texto complejo:** "El protocolo de comunicación asíncrona permite la transmisión de datos sin sincronización temporal."
            **Texto simple:** "Este método permite enviar información sin esperar a que termine el envío anterior."
            """)
    
    return demo

# Crear y lanzar la aplicación
if __name__ == "__main__":
    print("Iniciando aplicación de asistente de texto y código...")
    
    # Crear la interfaz
    app = create_autocomplete_interface()
    
    # Lanzar la aplicación
    app.launch(
        share=False,  # Cambiar a True si quieres compartir públicamente
        server_name="0.0.0.0",  # Permite acceso desde otras máquinas en la red local
        server_port=7860,  # Puerto por defecto de Gradio
        show_error=True,
        debug=False
    )