Spaces:
Sleeping
Sleeping
| import os | |
| import gradio as gr | |
| import torch | |
| from v1.usta_model import UstaModel | |
| from v1.usta_tokenizer import UstaTokenizer | |
| # Load the model and tokenizer | |
| def load_model(): | |
| try: | |
| u_tokenizer = UstaTokenizer("v1/tokenizer.json") | |
| print("✅ Tokenizer loaded successfully! vocab size:", len(u_tokenizer.vocab)) | |
| # Model parameters - adjust these to match your trained model | |
| context_length = 32 | |
| vocab_size = len(u_tokenizer.vocab) | |
| embedding_dim = 12 | |
| num_heads = 4 | |
| num_layers = 8 | |
| # Load the model | |
| u_model = UstaModel( | |
| vocab_size=vocab_size, | |
| embedding_dim=embedding_dim, | |
| num_heads=num_heads, | |
| context_length=context_length, | |
| num_layers=num_layers | |
| ) | |
| # Load the trained weights if available | |
| model_path = "v1/u_model.pth" | |
| if not os.path.exists(model_path): | |
| print("❌ Model file not found at", model_path) | |
| # Download the model file from GitHub | |
| try: | |
| print("📥 Downloading model weights from GitHub...") | |
| import requests | |
| url = "https://github.com/malibayram/llm-from-scratch/raw/main/u_model_4000.pth" | |
| headers = { | |
| 'Accept': 'application/octet-stream', | |
| 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36' | |
| } | |
| response = requests.get(url, headers=headers) | |
| response.raise_for_status() # Raise an exception for bad status codes | |
| # Check if we got a proper binary file (PyTorch files start with specific bytes) | |
| if response.content[:4] != b'PK\x03\x04' and b'<html' in response.content[:100].lower(): | |
| raise Exception("Downloaded HTML instead of binary file - check URL") | |
| print(f"📦 Downloaded {len(response.content)} bytes") | |
| # Create v1 directory if it doesn't exist | |
| os.makedirs("v1", exist_ok=True) | |
| # Save the model weights to the local file system | |
| with open(model_path, "wb") as f: | |
| f.write(response.content) | |
| print("✅ Model weights saved successfully!") | |
| except Exception as e: | |
| print(f"❌ Failed to download model weights: {e}") | |
| print("Using random initialization.") | |
| if os.path.exists(model_path): | |
| try: | |
| u_model.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=False)) | |
| u_model.eval() | |
| print("✅ Model weights loaded successfully!") | |
| except Exception as e: | |
| print(f"⚠️ Warning: Could not load trained weights: {e}") | |
| print("Using random initialization.") | |
| else: | |
| print(f"⚠️ Model file not found at {model_path}. Using random initialization.") | |
| return u_model, u_tokenizer | |
| except Exception as e: | |
| print(f"❌ Error loading model: {e}") | |
| raise e | |
| # Initialize model and tokenizer globally | |
| try: | |
| model, tokenizer = load_model() | |
| print("🚀 UstaModel and tokenizer initialized successfully!") | |
| except Exception as e: | |
| print(f"❌ Failed to initialize model: {e}") | |
| model, tokenizer = None, None | |
| def respond( | |
| message, | |
| history: list[tuple[str, str]], | |
| system_message, | |
| max_tokens, | |
| temperature, | |
| top_p, | |
| ): | |
| """ | |
| Generate a response using the UstaModel | |
| """ | |
| if model is None or tokenizer is None: | |
| yield "Sorry, the UstaModel is not available. Please try again later." | |
| return | |
| try: | |
| # For UstaModel, we'll use the message directly (ignoring system_message for now) | |
| # since it's a simpler model focused on geographical knowledge | |
| # Encode the input message | |
| tokens = tokenizer.encode(message) | |
| # Make sure we don't exceed context length | |
| if len(tokens) > 25: # Leave some room for generation | |
| tokens = tokens[-25:] | |
| # Generate response | |
| with torch.no_grad(): | |
| # Use max_tokens parameter, but cap it at reasonable limit for this model | |
| actual_max_tokens = min(max_tokens, 32 - len(tokens)) | |
| generated_tokens = model.generate(tokens, actual_max_tokens) | |
| # Decode the generated tokens | |
| response = tokenizer.decode(generated_tokens) | |
| # Clean up the response (remove the original input) | |
| original_text = tokenizer.decode(tokens.tolist()) | |
| if response.startswith(original_text): | |
| response = response[len(original_text):] | |
| # Clean up any unwanted tokens | |
| response = response.replace("<unk>", "").replace("<pad>", "").strip() | |
| if not response: | |
| response = "I'm not sure how to respond to that with my geographical knowledge." | |
| # Yield the response (to maintain compatibility with streaming interface) | |
| yield response | |
| except Exception as e: | |
| yield f"Sorry, I encountered an error: {str(e)}" | |
| """ | |
| For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface | |
| """ | |
| demo = gr.ChatInterface( | |
| respond, | |
| additional_inputs=[ | |
| gr.Textbox( | |
| value="You are Usta, a geographical knowledge assistant trained from scratch.", | |
| label="System message", | |
| info="Note: This model focuses on geographical knowledge (countries, capitals, cities)" | |
| ), | |
| gr.Slider(minimum=1, maximum=30, value=20, step=1, label="Max new tokens"), | |
| gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature"), | |
| gr.Slider( | |
| minimum=0.1, | |
| maximum=1.0, | |
| value=0.95, | |
| step=0.05, | |
| label="Top-p (nucleus sampling)", | |
| info="Note: This parameter is not used by UstaModel but kept for interface compatibility" | |
| ), | |
| ], | |
| title="🤖 Usta Model Chat", | |
| description="Chat with a custom transformer language model built from scratch! This model specializes in geographical knowledge including countries, capitals, and cities." | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |