equ1/generative_python_transformer

1from transformers import AutoTokenizer, AutoModelWithLMHead
2import gradio as gr
3
4# inference function
5def inference(inp):
6    tokenizer = AutoTokenizer.from_pretrained("./")
7    model = AutoModelWithLMHead.from_pretrained("./")
8
9    input_ids = tokenizer.encode(inp, return_tensors="pt")
10    beam_output = model.generate(input_ids, 
11                               max_length=512,
12                               num_beams=10,
13                               temperature=0.7,
14                               no_repeat_ngram_size=5,
15                               num_return_sequences=1,
16                               )
17  
18    output = []
19    for beam in beam_output:
20        out = tokenizer.decode(beam)
21        fout = out.replace("<N>", "\n")
22        output.append(fout)
23
24    return '\n'.join(output)
25
26desc = """
27        Enter some Python code and click submit to see the model's autocompletion.\n
28        
29        Best results have been observed with the prompt of \"import\".\n
30
31        Please note that outputs are reflective of a model trained on a measly 40 MBs of text data for 
32        a single epoch of ~16 GPU hours. Given more data and training time, the autocompletion should be much stronger.\n
33        
34        Computation will take some time.
35        """
36
37# Creates and launches gradio interface
38gr.Interface(fn=inference,
39            inputs=gr.inputs.Textbox(lines=5, label="Input Text"),
40            outputs=gr.outputs.Textbox(),
41            title="Generative Python Transformer",
42            description=desc,
43            ).launch()
44