Update app.py

app.py

@@ -1,409 +1,147 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PretrainedConfig
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import math
 import os
-import sys
-import transformers
-from dotenv import load_dotenv
-from huggingface_hub import login
-
-class RMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-5):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.eps = eps
-
-    def forward(self, x):
-        variance = x.pow(2).mean(-1, keepdim=True)
-        x = x * torch.rsqrt(variance + self.eps)
-        return self.weight * x
-
-class LlamaAttention(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
-        self.head_dim = config.hidden_size // config.num_attention_heads
-        
-        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
-        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
-        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
-        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
-
-    def forward(self, hidden_states, attention_mask=None):
-        batch_size, seq_length, _ = hidden_states.size()
-        
-        q = self.q_proj(hidden_states).view(batch_size, seq_length, self.num_heads, self.head_dim)
-        k = self.k_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
-        v = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
-        
-        if self.num_kv_heads < self.num_heads:
-            k = k.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
-            v = v.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
-        
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-        
-        attention_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
-        
-        if attention_mask is not None:
-            attention_scores = attention_scores + attention_mask
-            
-        attention_probs = F.softmax(attention_scores, dim=-1)
-        context = torch.matmul(attention_probs, v)
-        
-        context = context.transpose(1, 2).contiguous()
-        context = context.view(batch_size, seq_length, -1)
-        
-        return self.o_proj(context)
-
-class LlamaMLP(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
-        self.act_fn = nn.SiLU()
-
-    def forward(self, x):
-        gate = self.act_fn(self.gate_proj(x))
-        up = self.up_proj(x)
-        return self.down_proj(gate * up)
-
-class LlamaDecoderLayer(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.self_attn = LlamaAttention(config)
-        self.mlp = LlamaMLP(config)
-        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-
-    def forward(self, hidden_states, attention_mask=None):
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.self_attn(hidden_states, attention_mask)
-        hidden_states = residual + hidden_states
-        
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-        
-        return hidden_states
-
-class SmolLM2Config(PretrainedConfig):
-    model_type = "smollm2"
+import torch
+import gradio as gr
+from train_optimized import GPT, GPTConfig
+from huggingface_hub import hf_hub_download
+import json
+
+# Cache for model and tokenizer
+MODEL = None
+CHARS = None
+STOI = None
+ITOS = None
+
+def initialize():
+    global MODEL, CHARS, STOI, ITOS
     
-    def __init__(
-        self,
-        vocab_size=49152,
-        hidden_size=576,
-        intermediate_size=1536,
-        num_hidden_layers=30,
-        num_attention_heads=9,
-        num_key_value_heads=3,
-        hidden_act="silu",
-        max_position_embeddings=2048,
-        initializer_range=0.041666666666666664,
-        rms_norm_eps=1e-5,
-        use_cache=True,
-        pad_token_id=None,
-        bos_token_id=0,
-        eos_token_id=0,
-        tie_word_embeddings=True,
-        rope_theta=10000.0,
-        **kwargs
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs
+    if MODEL is None:
+        print("Loading model and tokenizer...")
+        # Download model files from HF Hub
+        config_path = hf_hub_download(repo_id="jatingocodeo/shakespeare-decoder", filename="config.json")
+        model_path = hf_hub_download(repo_id="jatingocodeo/shakespeare-decoder", filename="pytorch_model.bin")
+        
+        # Load config
+        with open(config_path, 'r') as f:
+            config_dict = json.load(f)
+        
+        # Initialize model with config
+        config = GPTConfig(
+            vocab_size=config_dict['vocab_size'],
+            n_layer=config_dict['n_layer'],
+            n_head=config_dict['n_head'],
+            n_embd=config_dict['n_embd'],
+            block_size=config_dict['block_size'],
+            dropout=config_dict['dropout'],
+            bias=config_dict['bias']
         )
-
-class SmolLM2ForCausalLM(PreTrainedModel):
-    config_class = SmolLM2Config
-    _no_split_modules = ["LlamaDecoderLayer"]
-    
-    def __init__(self, config):
-        super().__init__(config)
-        self.config = config
-        
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
-        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        model = GPT(config)
         
-        if config.tie_word_embeddings:
-            self.lm_head.weight = self.embed_tokens.weight
-            
-    def forward(self, input_ids, attention_mask=None, labels=None):
-        hidden_states = self.embed_tokens(input_ids)
-        
-        # Create causal attention mask if none provided
-        if attention_mask is None:
-            attention_mask = torch.triu(
-                torch.ones((input_ids.size(1), input_ids.size(1)), dtype=torch.bool, device=input_ids.device),
-                diagonal=1
-            )
-            attention_mask = attention_mask.unsqueeze(0).unsqueeze(0)
-            attention_mask = attention_mask * -1e4
-        
-        for layer in self.layers:
-            hidden_states = layer(hidden_states, attention_mask)
-            
-        hidden_states = self.norm(hidden_states)
-        logits = self.lm_head(hidden_states)
-        
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
-            
-        return logits if loss is None else (loss, logits)
-
-    def prepare_inputs_for_generation(self, input_ids, **kwargs):
-        return {
-            "input_ids": input_ids,
-            "attention_mask": kwargs.get("attention_mask", None)
-        }
-
-# Register the model architecture
-from transformers import AutoConfig, AutoModelForCausalLM
-AutoConfig.register("smollm2", SmolLM2Config)
-AutoModelForCausalLM.register(SmolLM2Config, SmolLM2ForCausalLM)
-
-# Load model and tokenizer
-model_id = "jatingocodeo/SmolLM2"
-
-def check_huggingface_access():
-    try:
-        from huggingface_hub import HfApi
-        api = HfApi()
+        # Load model weights
+        state_dict = torch.load(model_path, map_location='cpu')
+        model.load_state_dict(state_dict)
+        model.eval()
+        MODEL = model
         
-        # Check if token exists
+        # Initialize tokenizer
         try:
-            token = os.getenv("HF_TOKEN")
-            if not token:
-                print("× Warning: HF_TOKEN environment variable not found")
-                print("  Please ensure you have set your HuggingFace token")
-                return False
-        except Exception as e:
-            print(f"× Error checking HF_TOKEN: {str(e)}")
-            return False
-        
-        # Check repository access
-        try:
-            print(f"Checking access to repository: {model_id}")
-            repo_info = api.repo_info(model_id, token=token)
-            print("✓ Repository access confirmed")
-            print(f"Repository info: {repo_info}")
-            return True
-        except Exception as e:
-            print(f"× Error accessing repository: {str(e)}")
-            return False
-    except Exception as e:
-        print(f"× Error in HuggingFace access check: {str(e)}")
-        return False
-
-def load_model():
+            # First try to download input.txt from the repository
+            input_path = hf_hub_download(repo_id="jatingocodeo/shakespeare-decoder", filename="input.txt")
+            with open(input_path, 'r', encoding='utf-8') as f:
+                text = f.read()
+        except:
+            print("Warning: Could not load input.txt from repository. Using default character set.")
+            # Use a comprehensive set of characters that would be in Shakespeare's text
+            text = """
+            ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789.,!?-_:;'"()[]{}\n </>\t@#$%^&*+=~`|\\
+            """
+        
+        CHARS = sorted(list(set(text)))
+        STOI = {ch:i for i,ch in enumerate(CHARS)}
+        ITOS = {i:ch for i,ch in enumerate(CHARS)}
+        
+        print(f"Model loaded successfully! Vocabulary size: {len(CHARS)} characters")
+        print("Available characters:", ''.join(CHARS))
+
+def generate_text(
+    prompt, 
+    max_new_tokens=100, 
+    temperature=0.8, 
+    top_k=50
+):
+    # Initialize if not already done
+    if MODEL is None:
+        initialize()
+    
+    # Encode the prompt
+    encode = lambda s: [STOI[c] for c in s]
+    decode = lambda l: ''.join([ITOS[i] for i in l])
+    
     try:
-        print("\n=== Starting model loading process ===")
-        print(f"Model ID: {model_id}")
-        
-        # Load environment variables
-        load_dotenv()
-        
-        # Get HF token
-        hf_token = os.getenv('HF_TOKEN')
-        if not hf_token:
-            raise ValueError("HF_TOKEN not found in environment variables")
-            
-        # Login to Hugging Face
-        login(hf_token)
-        
-        print("\n1. Loading tokenizer...")
-        try:
-            tokenizer = AutoTokenizer.from_pretrained(
-                model_id,
-                use_auth_token=hf_token,
-                trust_remote_code=True
-            )
-            print("✓ Tokenizer loaded successfully")
-            print(f"Tokenizer type: {type(tokenizer)}")
-            print(f"Vocabulary size: {len(tokenizer)}")
-        except Exception as e:
-            print(f"× Error loading tokenizer: {str(e)}")
-            raise
-        
-        print("\n2. Adding special tokens...")
-        try:
-            special_tokens = {
-                'pad_token': '[PAD]',
-                'eos_token': '</s>',
-                'bos_token': '<s>'
-            }
-            num_added = tokenizer.add_special_tokens(special_tokens)
-            print(f"✓ Added {num_added} special tokens")
-            print(f"Special tokens: {tokenizer.special_tokens_map}")
-        except Exception as e:
-            print(f"× Error adding special tokens: {str(e)}")
-            raise
-        
-        print("\n3. Loading model...")
-        try:
-            model = AutoModelForCausalLM.from_pretrained(
-                model_id,
-                use_auth_token=hf_token,
-                trust_remote_code=True,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                low_cpu_mem_usage=True
-            )
-            print("✓ Model loaded successfully")
-            print(f"Model type: {type(model)}")
-        except Exception as e:
-            print(f"× Error loading model: {str(e)}")
-            raise
-        
-        print("\n4. Moving model to device...")
-        try:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            print(f"Selected device: {device}")
-            model = model.to(device)
-            print("✓ Model moved to device successfully")
-        except Exception as e:
-            print(f"× Error moving model to device: {str(e)}")
-            raise
-        
-        print("\n=== Model loading completed successfully! ===")
-        return model, tokenizer
-        
+        # Convert prompt to tensor
+        x = torch.tensor(encode(prompt), dtype=torch.long)[None,...]
+        
+        # Generate
+        with torch.no_grad():
+            y = MODEL.generate(x, max_new_tokens, temperature, top_k)[0]
+        
+        # Decode and return
+        generated_text = decode(y.tolist())
+        return generated_text
+    except KeyError:
+        return "Error: The prompt contains characters that are not in the training data. Please use only standard English characters."
     except Exception as e:
-        print("\n!!! ERROR IN MODEL LOADING !!!")
-        print(f"Error type: {type(e).__name__}")
-        print(f"Error message: {str(e)}")
-        print("\nFull traceback:")
-        import traceback
-        traceback.print_exc()
-        print("\nAdditional debug info:")
-        print(f"Python version: {sys.version}")
-        print(f"PyTorch version: {torch.__version__}")
-        print(f"Transformers version: {transformers.__version__}")
-        print(f"CUDA available: {torch.cuda.is_available()}")
-        if torch.cuda.is_available():
-            print(f"CUDA version: {torch.version.cuda}")
-        print("\nEnvironment variables:")
-        print(f"HF_TOKEN set: {'HF_TOKEN' in os.environ}")
-        raise
+        return f"Error generating text: {str(e)}"
 
-def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
-    try:
-        print("\n=== Starting text generation ===")
-        print(f"Input prompt: {prompt}")
-        print(f"Parameters: max_length={max_length}, temperature={temperature}, top_k={top_k}")
-        
-        if not hasattr(generate_text, "model"):
-            print("\n1. First call - loading model...")
-            generate_text.model, generate_text.tokenizer = load_model()
-        
-        if not prompt.strip():
-            print("× Empty prompt received")
-            return "Please enter a prompt."
-        
-        print("\n2. Processing prompt...")
-        if not prompt.startswith(generate_text.tokenizer.bos_token):
-            prompt = generate_text.tokenizer.bos_token + prompt
-            print("Added BOS token to prompt")
-        
-        print("\n3. Encoding prompt...")
-        try:
-            input_ids = generate_text.tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=2048)
-            print(f"Encoded shape: {input_ids.shape}")
-            input_ids = input_ids.to(generate_text.model.device)
-            print("✓ Encoding successful")
-        except Exception as e:
-            print(f"× Error encoding prompt: {str(e)}")
-            raise
-        
-        print("\n4. Generating text...")
-        try:
-            with torch.no_grad():
-                output_ids = generate_text.model.generate(
-                    input_ids,
-                    max_length=min(max_length + len(input_ids[0]), 2048),
-                    temperature=temperature,
-                    top_k=top_k,
-                    do_sample=True,
-                    pad_token_id=generate_text.tokenizer.pad_token_id,
-                    eos_token_id=generate_text.tokenizer.eos_token_id,
-                    num_return_sequences=1
-                )
-            print(f"Generation shape: {output_ids.shape}")
-        except Exception as e:
-            print(f"× Error during generation: {str(e)}")
-            raise
-        
-        print("\n5. Decoding output...")
-        try:
-            generated_text = generate_text.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-            print("✓ Decoding successful")
-            print(f"Output length: {len(generated_text)}")
-        except Exception as e:
-            print(f"× Error decoding output: {str(e)}")
-            raise
-        
-        print("\n=== Generation completed successfully! ===")
-        return generated_text.strip()
-    
-    except Exception as e:
-        print("\n!!! ERROR IN TEXT GENERATION !!!")
-        print(f"Error type: {type(e).__name__}")
-        print(f"Error message: {str(e)}")
-        print("\nFull traceback:")
-        import traceback
-        traceback.print_exc()
-        return f"An error occurred: {str(e)}"
+# Initialize on startup
+initialize()
 
 # Create Gradio interface
-iface = gr.Interface(
+demo = gr.Interface(
     fn=generate_text,
     inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=2),
-        gr.Slider(minimum=10, maximum=200, value=100, step=1, label="Max Length"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top K"),
+        gr.Textbox(
+            label="Prompt",
+            placeholder="Enter your prompt here...",
+            lines=5
+        ),
+        gr.Slider(
+            label="Max New Tokens",
+            minimum=10,
+            maximum=500,
+            value=100,
+            step=10
+        ),
+        gr.Slider(
+            label="Temperature",
+            minimum=0.1,
+            maximum=2.0,
+            value=0.8,
+            step=0.1
+        ),
+        gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=100,
+            value=50,
+            step=1
+        ),
     ],
-    outputs=gr.Textbox(label="Generated Text", lines=5),
-    title="SmolLM2 Text Generator",
-    description="""Generate text using the fine-tuned SmolLM2 model.
-    - Max Length: Controls the length of generated text
-    - Temperature: Controls randomness (higher = more creative)
-    - Top K: Controls diversity of word choices""",
+    outputs=gr.Textbox(label="Generated Text", lines=10),
+    title="Shakespeare GPT",
+    description="""
+    This is a GPT model trained on Shakespeare's text. Enter a prompt and the model will continue it in Shakespeare's style.
+    
+    Parameters:
+    - Temperature: Higher values make the output more random, lower values make it more deterministic
+    - Top-k: Number of highest probability tokens to consider at each step
+    - Max New Tokens: Maximum number of tokens to generate
+    """,
     examples=[
-        ["Once upon a time", 100, 0.7, 50],
-        ["The quick brown fox", 150, 0.8, 40],
-        ["In a galaxy far far away", 200, 0.9, 30],
-    ],
-    allow_flagging="never"
+        ["To be, or not to be,", 100, 0.8, 50],
+        ["Friends, Romans, countrymen,", 150, 0.7, 40],
+        ["Now is the winter of", 200, 0.9, 30],
+    ]
 )
 
 if __name__ == "__main__":
-    print("Starting Gradio interface...")
-    iface.launch(debug=True) 
+    demo.launch()