Update app.py

app.py

@@ -3,18 +3,63 @@ import torch
 import torch.nn as nn
 from torch.nn import functional as F
 import json
-import os # <-- Added for file path checks
+import os
 
-# --- Model Definition (same as before) ---
-# NOTE: The model class MUST be defined in your app.py file
-# so that torch.load knows how to reconstruct it.
+# --- Model Hyperparameters (same as before) ---
 batch_size = 32
 block_size = 8
 n_embd = 32
 n_head = 4
 n_layer = 4
 dropout = 0.0
+max_iters = 3000
+eval_interval = 300
+learning_rate = 1e-2
+eval_iters = 200
+
+# --- Data Preparation & Vocabulary Creation ---
+file_path = 'dataset.jsonl'
+corpus = ""
+try:
+    with open(file_path, 'r') as f:
+        for line in f:
+            data_point = json.loads(line)
+            corpus += data_point['header'] + '\n' + data_point['formal_statement'] + '\n'
+except FileNotFoundError:
+    print(f"Error: The file '{file_path}' was not found.")
+    exit()
+except (json.JSONDecodeError, KeyError):
+    print(f"Error: There was a problem parsing a line in '{file_path}'. Check for malformed JSON or missing keys.")
+    exit()
 
+if not corpus:
+    print("Error: The corpus is empty.")
+    exit()
+
+chars = sorted(list(set(corpus)))
+vocab_size = len(chars)
+stoi = {ch: i for i, ch in enumerate(chars)}
+itos = {i: ch for i, ch in enumerate(chars)}
+encode = lambda s: [stoi.get(c, 0) for c in s]
+decode = lambda l: ''.join([itos[i] for i in l])
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Split the data for training and validation
+data = torch.tensor(encode(corpus), dtype=torch.long)
+n = int(0.9 * len(data))
+train_data = data[:n]
+val_data = data[n:]
+
+def get_batch(split):
+    data = train_data if split == 'train' else val_data
+    ix = torch.randint(len(data) - block_size, (batch_size,))
+    x = torch.stack([data[i:i + block_size] for i in ix])
+    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
+    x, y = x.to(device), y.to(device)
+    return x, y
+
+# --- Model Definition (same as before) ---
 class Head(nn.Module):
     def __init__(self, head_size):
         super().__init__()
@@ -81,7 +126,6 @@ class LanguageModel(nn.Module):
         self.lm_head = nn.Linear(n_embd, vocab_size)
         self.block_size = block_size
         self.vocab_size = vocab_size
-
     def forward(self, idx, targets=None):
         B, T = idx.shape
         tok_emb = self.token_embedding_table(idx)
@@ -97,7 +141,6 @@ class LanguageModel(nn.Module):
             targets = targets.view(B * T)
             loss = F.cross_entropy(logits, targets)
         return logits, loss
-
     def generate(self, idx, max_new_tokens):
         for _ in range(max_new_tokens):
             idx_cond = idx[:, -self.block_size:]
@@ -108,77 +151,76 @@ class LanguageModel(nn.Module):
             idx = torch.cat((idx, idx_next), dim=1)
         return idx
 
+# --- Training and Generation ---
+model = LanguageModel(vocab_size, block_size, n_embd, n_head, n_layer, dropout)
+m = model.to(device)
+
+# --- Check if a trained model exists, otherwise train a new one ---
+model_file = 'model.pt'
+if os.path.exists(model_file):
+    print(f"Loading existing model from {model_file}")
+    try:
+        model.load_state_dict(torch.load(model_file, map_location=device))
+    except RuntimeError as e:
+        print(f"Error loading model: {e}")
+        print("Model file might be incompatible with current vocabulary. Retraining...")
+        # If loading fails, fall through to training logic
+        model.train() # Set back to train mode just in case
+else:
+    print("No trained model found. Starting a new training session...")
+    
+    # Define a helper function for loss estimation
+    @torch.no_grad()
+    def estimate_loss():
+        out = {}
+        model.eval()
+        for split in ['train', 'val']:
+            losses = torch.zeros(eval_iters)
+            for k in range(eval_iters):
+                X, Y = get_batch(split)
+                logits, loss = model(X, Y)
+                losses[k] = loss.item()
+            out[split] = losses.mean()
+        model.train()
+        return out
 
-# --- Utility functions and model loading ---
-# Read from the same dataset file to ensure vocabulary consistency
-file_path = 'dataset.jsonl'
-corpus = ""
-try:
-    with open(file_path, 'r') as f:
-        for line in f:
-            data_point = json.loads(line)
-            corpus += data_point['header'] + '\n' + data_point['formal_statement'] + '\n'
-except FileNotFoundError:
-    print(f"Error: The file '{file_path}' was not found.")
-    # Exit or handle gracefully as the app can't run without the dataset
-    exit()
-except json.JSONDecodeError:
-    print(f"Error: There was a problem parsing a line in '{file_path}'.")
-    exit()
-except KeyError:
-    print(f"Error: A line in '{file_path}' does not have the expected keys.")
-    exit()
-
-if not corpus:
-    print("Error: The corpus is empty.")
-    exit()
-
-chars = sorted(list(set(corpus)))
-vocab_size = len(chars)
-stoi = {ch: i for i, ch in enumerate(chars)}
-itos = {i: ch for i, ch in enumerate(chars)}
-# Corrected the encode function
-encode = lambda s: [stoi[c] for c in s]
-decode = lambda l: ''.join([itos[i] for i in l])
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    # The training loop
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+    for iter in range(max_iters):
+        if iter % eval_interval == 0:
+            losses = estimate_loss()
+            print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
+        xb, yb = get_batch('train')
+        logits, loss = model(xb, yb)
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+    
+    torch.save(m.state_dict(), model_file)
+    print(f"Training complete. Model saved to {model_file}")
 
-# Load the trained model.
-model = LanguageModel(vocab_size, block_size, n_embd, n_head, n_layer, dropout)
-model.load_state_dict(torch.load('model.pt', map_location=device))
-model.eval()  # Set the model to evaluation mode for inference.
+model.eval()
 model.to(device)
 
-
 # --- Gradio UI & Inference function ---
 def generate_text_chat(message, history):
-    # We'll just use the most recent message as the prompt.
     prompt = message
-    # You can adjust this to a different number of tokens if you like.
     max_new_tokens = 50 
-
-    # Encode the prompt text into tokens.
     encoded_prompt = [stoi.get(c, 0) for c in prompt]
     if not encoded_prompt:
         return "Prompt is empty or contains unknown characters."
-
     context = torch.tensor(encoded_prompt, dtype=torch.long, device=device).unsqueeze(0)
-    # Generate new tokens.
     generated_text_indices = model.generate(context, max_new_tokens=max_new_tokens)
-    # Decode the tokens back into text.
     generated_text = decode(generated_text_indices[0].tolist())
-    
-    # Return only the newly generated part of the text, removing the original prompt
     return generated_text[len(prompt):]
 
-# Using gr.ChatInterface for a conversational experience
 demo = gr.ChatInterface(
     fn=generate_text_chat,
     title="Tiny Language Model Chat",
     description="A simple character-level language model trained in PyTorch, now with a chat interface.",
-    # You can customize these components further if you like
     chatbot=gr.Chatbot(height="500px"),
     textbox=gr.Textbox(placeholder="Ask me anything...", container=False, scale=7),
     theme="soft",
 )
 
-demo.launch()
+demo.launch()