Update app.py

app.py

@@ -2,25 +2,32 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# Model repo
 repo_id = "theguywhosucks/mochaV2"
 
-# Load the tokenizer from the repo (uses tokenizer.json internally)
+# Load the tokenizer shipped with the model (tokenizer.json internally)
 tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=False)
 
-# GPT2-style models often don't have a pad token
+# GPT2-style models often lack a pad token; set it to eos_token
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-# Load the model (safetensors used automatically)
+# Load model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModelForCausalLM.from_pretrained(
     repo_id,
-    dtype=torch.float32,  # torch_dtype is deprecated
-    trust_remote_code=True
+    trust_remote_code=True,  # required if model uses custom code
+    dtype=torch.float32      # torch_dtype is deprecated, use dtype
 )
 model.to(device)
 model.eval()
 
+# Optional: confirm vocab sizes match
+assert tokenizer.vocab_size == model.config.vocab_size, (
+    f"Tokenizer vocab size ({tokenizer.vocab_size}) does not match model ({model.config.vocab_size})"
+)
+
+# Gradio function
 def complete_sentence(prompt, max_new_tokens=50, temperature=0.7):
     # Tokenize input safely
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
@@ -34,6 +41,7 @@ def complete_sentence(prompt, max_new_tokens=50, temperature=0.7):
             pad_token_id=tokenizer.pad_token_id
         )
 
+    # Decode output, skipping special tokens
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Launch Gradio app