""" HALO-S Playground — Gradio App para HuggingFace Spaces. Carga modelos HALO-S desde HuggingFace Hub usando safetensors y permite generación interactiva de texto. Compatible con: - HALO-S v1.x (CharacterTokenizer, vocab_size=256) - HALO-S v2.x (tiktoken GPT-2 BPE, vocab_size=50257) """ import torch import gradio as gr try: import tiktoken _TIKTOKEN_AVAILABLE = True except ImportError: _TIKTOKEN_AVAILABLE = False try: from safetensors.torch import load_file as load_safetensors except ImportError: load_safetensors = None from huggingface_hub import hf_hub_download from halo import HaloConfig, HaloSModel, CharacterTokenizer from halo.core.device import optimize_for_device, get_optimal_device # ============================================================ # MODELOS DISPONIBLES # ============================================================ MODELS = { "HALO-S Usmall (3.7M)": { "repo": "BUEORM/HALO-S-Usmall", "checkpoint": "model.safetensors", "tokenizer": "char", }, "HALO-S V1 (20M)": { "repo": "BUEORM/HALO-S-V1", "checkpoint": "model.safetensors", "tokenizer": "char", }, "HALO-S V2 (70M)": { "repo": "BUEORM/HALO-S-large", "checkpoint": "model.safetensors", "tokenizer": "gpt2", }, } # ============================================================ # CONFIGURACIONES — deben coincidir EXACTAMENTE con las usadas al entrenar # ============================================================ CONFIGS = { "HALO-S Usmall (3.7M)": HaloConfig( vocab_size=256, hidden_size=256, num_layers=4, num_heads=8, num_kv_heads=2, num_globals=2, local_window=32, dilated_offsets=[1, 2, 4, 8, 16], num_random=2, dropout=0.0, max_seq_len=256, use_swiglu=False, ), "HALO-S V1 (20M)": HaloConfig( vocab_size=256, hidden_size=512, num_layers=6, num_heads=8, num_kv_heads=2, num_globals=2, local_window=64, dilated_offsets=[1, 2, 4, 8, 16, 32, 64], num_random=2, dropout=0.0, max_seq_len=1024, use_swiglu=False, ), "HALO-S V2 (70M)": HaloConfig( vocab_size=50257, hidden_size=512, num_layers=6, num_heads=8, num_kv_heads=2, num_globals=2, local_window=64, dilated_offsets=[1, 2, 4, 8, 16, 32, 64, 128], num_random=2, dropout=0.0, max_seq_len=1024, use_swiglu=False, ), } # ============================================================ # CACHE DE MODELOS Y TOKENIZERS # ============================================================ _model_cache = {} _tokenizer_cache = {} # ============================================================ # FUNCIONES DE CARGA # ============================================================ def get_tokenizer(model_name): """Obtiene el tokenizer correcto para cada modelo.""" if model_name in _tokenizer_cache: return _tokenizer_cache[model_name] tok_type = MODELS[model_name]["tokenizer"] if tok_type == "char": tokenizer = CharacterTokenizer() elif tok_type == "gpt2": if not _TIKTOKEN_AVAILABLE: raise ImportError("tiktoken no instalado. pip install tiktoken") tokenizer = tiktoken.get_encoding("gpt2") else: raise ValueError(f"Tokenizer desconocido: {tok_type}") _tokenizer_cache[model_name] = tokenizer return tokenizer def load_model(model_name): """Carga un modelo desde HuggingFace Hub con safetensors.""" if model_name in _model_cache: return _model_cache[model_name] info = MODELS[model_name] config = CONFIGS[model_name] # Descargar el archivo safetensors desde HF Hub checkpoint_path = hf_hub_download( repo_id=info["repo"], filename=info["checkpoint"], ) # Crear modelo con la configuración correcta model = HaloSModel(config) # Cargar pesos desde safetensors if load_safetensors is not None and checkpoint_path.endswith(".safetensors"): state_dict = load_safetensors(checkpoint_path, device="cpu") else: state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True) # strict=False para compatibilidad entre versiones (w3 de SwiGLU puede faltar) model.load_state_dict(state_dict, strict=False) model.eval() # Optimizar para el dispositivo actual (CPU en Spaces, GPU si disponible) # NO usar torch.compile — en generación autoregresiva la compilación # del grafo tarda más que la inferencia misma para modelos pequeños device = get_optimal_device() model = model.to(device) # Solo aplicar optimizaciones ligeras (TF32, threads) sin compile import os if "cuda" in device: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True else: torch.set_num_threads(os.cpu_count() or 4) _model_cache[model_name] = model return model # ============================================================ # GENERACIÓN # ============================================================ def generate_text(prompt, model_name, max_new_tokens, temperature): """Genera texto con el modelo seleccionado.""" if not prompt or not prompt.strip(): return "⚠️ Por favor ingresa un prompt." try: model = load_model(model_name) tokenizer = get_tokenizer(model_name) device = next(model.parameters()).device # Tokenizar el prompt tok_type = MODELS[model_name]["tokenizer"] if tok_type == "char": token_ids = tokenizer.encode(prompt) else: token_ids = tokenizer.encode(prompt) if len(token_ids) == 0: return "⚠️ El prompt no produjo tokens válidos." # Limitar longitud del prompt al max_seq_len del modelo config = CONFIGS[model_name] max_ctx = config.max_seq_len - config.num_globals - int(max_new_tokens) if max_ctx < 1: max_ctx = config.max_seq_len // 2 token_ids = token_ids[-max_ctx:] input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) # Generar with torch.no_grad(): output = model.generate( input_ids=input_ids, max_new_tokens=int(max_new_tokens), temperature=float(temperature), top_k=50, top_p=0.9, ) # Decodificar output_ids = output[0].tolist() if tok_type == "char": text = tokenizer.decode(output_ids) else: text = tokenizer.decode(output_ids) return text except Exception as e: return f"❌ Error durante la generación:\n\n{type(e).__name__}: {str(e)}" # ============================================================ # INTERFAZ GRADIO # ============================================================ DESCRIPTION = """ # 🌀 HALO-S Playground **Generación de texto con modelos HALO-S** — Atención dispersa O(N×K) ### Modelos disponibles: | Modelo | Parámetros | Tokenizer | Entrenado en | |--------|-----------|-----------|--------------| | HALO-S Usmall | 3.7M | Character (256) | WikiText-2 | | HALO-S V1 | 20M | Character (256) | WikiText-103 | | HALO-S V2 | 70M | GPT-2 BPE (50257) | WikiText-103 | ### Notas: - Los modelos son experimentales y fueron entrenados con datos limitados - La calidad de generación depende del modelo y los hiperparámetros - Para V2 usa prompts en inglés (entrenado con BPE en texto inglés) - Compatible con HALO-S v1.x y v2.x [📦 PyPI](https://pypi.org/project/pyhalos/) | [💻 GitHub](https://github.com/bueormnew/pyhalo) | [📖 Docs](https://github.com/bueormnew/pyhalo/tree/main/docs) """ demo = gr.Interface( fn=generate_text, inputs=[ gr.Textbox( lines=5, label="📝 Prompt", placeholder="The history of artificial intelligence...", value="The history of", ), gr.Dropdown( choices=list(MODELS.keys()), value="HALO-S Usmall (3.7M)", label="🧠 Modelo", ), gr.Slider( minimum=10, maximum=512, value=100, step=10, label="📏 Max New Tokens", ), gr.Slider( minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="🌡️ Temperature", ), ], outputs=gr.Textbox( lines=15, label="✨ Texto Generado", ), title="🌀 HALO-S Playground", description=DESCRIPTION, examples=[ ["The history of", "HALO-S Usmall (3.7M)", 100, 0.8], ["Machine learning is", "HALO-S V1 (20M)", 150, 0.7], ["The United States of America", "HALO-S V2 (70M)", 200, 0.8], ], cache_examples=False, ) # ============================================================ # LANZAR # ============================================================ if __name__ == "__main__": demo.launch()