Add CodeT5 model integration

app.py

@@ -1,7 +1,49 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-def greet(name):
-    return "Hello " + name + "!!"
+# 加载 CodeT5 模型
+model_name = "Salesforce/codet5-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+def generate_code(prompt: str, max_length: int = 128) -> str:
+    """代码生成/补全"""
+    if not prompt.strip():
+        return ""
+    
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_length=max_length,
+            num_beams=4,
+            early_stopping=True
+        )
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# 创建 Gradio 界面
+demo = gr.Interface(
+    fn=generate_code,
+    inputs=[
+        gr.Textbox(
+            label="Prompt",
+            placeholder="输入代码描述或代码片段，例如：def fibonacci(n):",
+            lines=5
+        ),
+        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 的代码生成模型。支持代码补全、代码生成等任务。",
+    examples=[
+        ["def fibonacci(n):", 128],
+        ["# Python function to calculate factorial", 128],
+        ["def quick_sort(arr):", 128],
+    ]
+)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+transformers>=4.30.0
+torch>=2.0.0