Switch to CodeGen-350M for better tokenizer compatibility

app.py

@@ -1,11 +1,11 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# 加载 CodeT5+ 模型（兼容性更好）
-model_name = "Salesforce/codet5p-220m"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
+# 加载 CodeGen 模型（GPT-2架构，兼容性好）
+model_name = "Salesforce/codegen-350M-mono"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 def generate_code(prompt: str, max_length: int = 128) -> str:
     """代码生成/补全"""
@@ -19,7 +19,8 @@ def generate_code(prompt: str, max_length: int = 128) -> str:
             **inputs,
             max_length=max_length,
             num_beams=4,
-            early_stopping=True
+            early_stopping=True,
+            pad_token_id=tokenizer.eos_token_id
         )
     
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -36,8 +37,8 @@ demo = gr.Interface(
         gr.Slider(32, 512, value=128, step=32, label="Max Length")
     ],
     outputs=gr.Textbox(label="Generated Code", lines=10),
-    title="CodeT5+ Code Generation",
-    description="基于 Salesforce CodeT5+ (220M) 的代码生成模型。支持代码补全、代码生成等任务。",
+    title="CodeGen Code Generation",
+    description="基于 Salesforce CodeGen (350M) 的代码生成模型。支持代码补全、代码生成等任务。",
     examples=[
         ["def fibonacci(n):", 128],
         ["# Python function to calculate factorial", 128],
@@ -46,4 +47,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()