File size: 861 Bytes
7d465ab
 
 
 
 
d3c2d2f
7d465ab
 
 
 
 
 
 
8c5dd78
7d465ab
 
 
 
 
 
626e0b9
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

model_name = "Salesforce/codegen-350M-mono"
codegen_token = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def codegen(intent):
    """Give input as text which reflects intent of the program.
    """
    #text = "Write a function which takes 2 numbers as input and returns the larger of the two."

    input_ids = codegen_token(intent, return_tensors="pt").input_ids
    outcode_ids = model.generate(input_ids, max_length=256)
    response = codegen_token.decode(outcode_ids[0], skip_special_tokens=True)
    return response

# UX
in_text = gr.Textbox(lines=1, label="Place your intent here.")
out = gr.Textbox(lines=1, label="Generated python code", placeholder="")
gr.Interface(codegen, inputs=in_text, outputs=out).launch()