Update app.py

app.py

@@ -5,29 +5,38 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # HF repo containing your model (with safetensors)
 repo_id = "theguywhosucks/mochaV2"
 
-# Load tokenizer from HF (no manual itos/stoi)
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=False)
 
-# Load model (safetensors will be used automatically if available)
+# GPT2-style models often don't have a pad token, set it to eos
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Load model (safetensors automatically used if available)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModelForCausalLM.from_pretrained(
     repo_id,
-    torch_dtype=torch.float32,  # or torch.float16 for faster GPU inference
+    dtype=torch.float32,  # torch_dtype is deprecated; use dtype
     trust_remote_code=True
 )
 model.to(device)
 model.eval()
 
-# Gradio function
+# Gradio completion function
 def complete_sentence(prompt, max_new_tokens=50, temperature=0.7):
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    # Encode input with proper padding
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(device)
+
     with torch.no_grad():
         outputs = model.generate(
-            input_ids,
+            **inputs,
             max_new_tokens=max_new_tokens,
             do_sample=True,
-            temperature=temperature
+            temperature=temperature,
+            pad_token_id=tokenizer.pad_token_id  # ensures safe embedding lookup
         )
+
+    # Decode output safely
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Launch Gradio app