How i build a UI demo like mentioned in readme (input to output )

[prameelareddi]

How i build a UI demo like mentioned in the readme (input to output )

https://github.com/salesforce/CodeGen/blob/main/assets/two.gif

[AndreasMeyerSFDC]

@prameela1610 the easiest way is to use gradio. Here is a code snippet implementing the demo:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen2-7B")
model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen2-7B", trust_remote_code=True, revision="main")

def generate_code(input_text):
    inputs = tokenizer(input_text, return_tensors="pt")
    sample = model.generate(**inputs, max_length=128)
    generated_code = tokenizer.decode(sample[0], truncate_before_pattern=[r"\n\n^#", "^'''", "\n\n\n"])
    return generated_code

inputs = gr.inputs.Textbox(label="Input Text")
outputs = gr.outputs.Textbox(label="Generated Code")
interface = gr.Interface(fn=generate_code, inputs=inputs, outputs=outputs)

interface.launch()

pip install gradio will do the trick ;)