Switch to CodeT5+ (codet5p-220m) model

app.py

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