Update app.py

app.py

@@ -1,32 +1,29 @@
 import gradio as gr
 import torch
-import accelerate
-# Added GPT2Config and GPT2LMHeadModel to the imports
 from transformers import AutoTokenizer, GenerationConfig, GPT2Config, GPT2LMHeadModel
 
 # 1. Load the standard GPT-2 tokenizer
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
 
-# 2. Define the configuration and load the model
-# We use GPT2Config to ensure the 'model_type' is explicitly handled
-config = GPT2Config.from_pretrained("gpt2") 
+# 2. Define the configuration with YOUR specific vocab size
+# Changed vocab_size to 50000 based on your error report
+config = GPT2Config.from_pretrained("gpt2", vocab_size=50000) 
 
-# Load your local .safetensors weights into the GPT2 structure
+# 3. Load the model and ignore the size mismatch warning
 model = GPT2LMHeadModel.from_pretrained(
     ".", 
     config=config,
     local_files_only=True,
     torch_dtype=torch.float32,
-    device_map="auto"
+    device_map="auto",
+    ignore_mismatched_sizes=True  # This is the key fix!
 )
 
-# 3. Load your specific generation settings from your local file
-# Ensure 'generation_config.json' is actually in your Space root folder
+# 4. Load your generation settings
 gen_config = GenerationConfig.from_pretrained(".", "generation_config.json")
 
 def generate_code(prompt):
-    # Move inputs to the same device as the model (CPU or GPU)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
     with torch.no_grad():
@@ -38,7 +35,7 @@ def generate_code(prompt):
     
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# 4. Gradio Interface
+# 5. Gradio Interface
 demo = gr.Interface(
     fn=generate_code,
     inputs=gr.Textbox(placeholder="Write a function to...", label="Input Prompt"),