import gradio as gr
from transformers import pipeline

# 加载 CodeGen 模型
code_generator = pipeline('text-generation', model='Salesforce/codegen-350M-multi')  # 确保模型名称正确

def generate_code(prompt):
    # 设置合理的 max_length
    result = code_generator(prompt, max_length=150, truncation=True)  # 使用合理的 max_length
    return result[0]['generated_text']

# 创建 Gradio 接口
iface = gr.Interface(fn=generate_code, inputs="text", outputs="text", title="Code Generator",
                     description="输入您的代码提示，生成代码片段。")
iface.launch(share=True)  # 设置为 True 以创建公共链接