# 加载模型和分词器
checkpoint = "bigcode/santacoder"
device = "cuda"  # 若要使用GPU，请将其设置为"cuda"，若使用CPU，请设置为"cpu"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint, revision="main", trust_remote_code=True).to(device)

# 创建代码生成函数
def generate_code(nl_description):
    inputs = tokenizer.encode(nl_description, return_tensors="pt").to(device)
    outputs = model.generate(inputs)
    return tokenizer.decode(outputs[0])

# 使用Gradio创建界面
input_interface = gr.inputs.Textbox(lines=5, placeholder="请输入自然语言描述...")
output_interface = gr.outputs.Code()

gr.Interface(fn=generate_code, inputs=input_interface, outputs=output_interface, title="SantaCoder代码生成器").launch()