import gradio as gr import torch from beeper_model import BeeperRoseGPT, generate from tokenizers import Tokenizer from huggingface_hub import hf_hub_download from safetensors.torch import load_file as load_safetensors # ---------------------------- # 🔧 Model versions configuration # ---------------------------- MODEL_VERSIONS = { "Beeper v1 (Original)": { "repo_id": "AbstractPhil/beeper-rose-tinystories-6l-512d-ctx512", "model_file": "beeper_rose_final.safetensors", "description": "Original Beeper trained on TinyStories" }, "Beeper v2 (Extended)": { "repo_id": "AbstractPhil/beeper-rose-v2", "model_file": "beeper_rose_final.safetensors", "description": "Beeper v2 with extended training (~15 epochs) on a good starting corpus of general knowledge." } } # Base configuration config = { "context": 512, "vocab_size": 8192, "dim": 512, "n_heads": 8, "n_layers": 6, "mlp_ratio": 4.0, "temperature": 0.9, "top_k": 40, "top_p": 0.9, "repetition_penalty": 1.1, "presence_penalty": 0.6, "frequency_penalty": 0.0, "resid_dropout": 0.1, "dropout": 0.0, "grad_checkpoint": False, "tokenizer_path": "beeper.tokenizer.json" } device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Global model and tokenizer variables infer = None tok = None current_version = None def load_model_version(version_name): """Load the selected model version""" global infer, tok, current_version if current_version == version_name and infer is not None: return f"Already loaded: {version_name}" version_info = MODEL_VERSIONS[version_name] try: # Download model and tokenizer files model_file = hf_hub_download( repo_id=version_info["repo_id"], filename=version_info["model_file"] ) tokenizer_file = hf_hub_download( repo_id=version_info["repo_id"], filename="tokenizer.json" ) # Initialize model infer = BeeperRoseGPT(config).to(device) # Load safetensors state_dict = load_safetensors(model_file, device=str(device)) infer.load_state_dict(state_dict) infer.eval() # Load tokenizer tok = Tokenizer.from_file(tokenizer_file) current_version = version_name return f"Successfully loaded: {version_name}" except Exception as e: return f"Error loading {version_name}: {str(e)}" # Load default model on startup load_status = load_model_version("Beeper v1 (Original)") print(load_status) # ---------------------------- # 💬 Gradio Chat Wrapper # ---------------------------- def beeper_reply(message, history, model_version, temperature=None, top_k=None, top_p=None): global infer, tok, current_version # Load model if version changed if model_version != current_version: status = load_model_version(model_version) if "Error" in status: return f"⚠️ {status}" # Check if model is loaded if infer is None or tok is None: return "⚠️ Model not loaded. Please select a version and try again." # Use defaults if not provided (for examples caching) if temperature is None: temperature = 0.9 if top_k is None: top_k = 40 if top_p is None: top_p = 0.9 # Build conversation context prompt_parts = [] if history: for h in history: if h[0]: # User message exists prompt_parts.append(f"User: {h[0]}") if h[1]: # Assistant response exists prompt_parts.append(f"Beeper: {h[1]}") # Add current message prompt_parts.append(f"User: {message}") prompt_parts.append("Beeper:") prompt = "\n".join(prompt_parts) # Generate response response = generate( model=infer, tok=tok, cfg=config, prompt=prompt, max_new_tokens=128, temperature=float(temperature), top_k=int(top_k), top_p=float(top_p), repetition_penalty=config["repetition_penalty"], presence_penalty=config["presence_penalty"], frequency_penalty=config["frequency_penalty"], device=device, detokenize=True ) # Clean up response - remove the prompt part if it's included if response.startswith(prompt): response = response[len(prompt):].strip() return response # ---------------------------- # 🖼️ Interface # ---------------------------- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown( """ # 🤖 Beeper - A Rose-based Tiny Language Model Hello! I'm Beeper, a small language model trained with love and care. Please be patient with me - I'm still learning! 💕 """ ) with gr.Row(): with gr.Column(scale=3): model_dropdown = gr.Dropdown( choices=list(MODEL_VERSIONS.keys()), value="Beeper v1 (Original)", label="Select Beeper Version", info="Choose which version of Beeper to chat with" ) with gr.Column(scale=7): version_info = gr.Markdown("**Current:** Beeper v1 - Original training on TinyStories") # Update version info when dropdown changes def update_version_info(version_name): info = MODEL_VERSIONS[version_name]["description"] return f"**Current:** {info}" model_dropdown.change( fn=update_version_info, inputs=[model_dropdown], outputs=[version_info] ) # Chat interface chatbot = gr.Chatbot(label="Chat with Beeper", type="messages", height=400) msg = gr.Textbox(label="Message", placeholder="Type your message here...") with gr.Row(): with gr.Column(scale=2): temperature_slider = gr.Slider(0.1, 1.5, value=0.9, step=0.1, label="Temperature") with gr.Column(scale=2): top_k_slider = gr.Slider(1, 100, value=40, step=1, label="Top-k") with gr.Column(scale=2): top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p") with gr.Row(): submit = gr.Button("Send", variant="primary") clear = gr.Button("Clear") # Examples gr.Examples( examples=[ ["Hello Beeper! How are you today?"], ["Can you tell me a story about a robot?"], ["What do you like to do for fun?"], ["What makes you happy?"], ["Tell me about your dreams"], ], inputs=msg ) # Handle chat def respond(message, chat_history, model_version, temperature, top_k, top_p): response = beeper_reply(message, chat_history, model_version, temperature, top_k, top_p) chat_history.append([message, response]) return "", chat_history msg.submit( respond, [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider], [msg, chatbot] ) submit.click( respond, [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider], [msg, chatbot] ) clear.click(lambda: None, None, chatbot, queue=False) if __name__ == "__main__": demo.launch()