go

app.py

@@ -14,7 +14,7 @@ huggingface_hub.login(token=hf_key)
 
 tokenizer = AutoTokenizer.from_pretrained(
     "bigcode/starcoderbase-3b"
-    #, use_auth_token=hf_key
+    # , use_auth_token=hf_key
 )
 vardecoder_model = AutoModelForCausalLM.from_pretrained(
     "ejschwartz/resym-vardecoder", torch_dtype=torch.bfloat16, device_map="auto"
@@ -40,18 +40,36 @@ example = """{
 @spaces.GPU
 def infer(input):
     line = json.loads(input)
-    first_token = line['output'].split(':')[0]
-    prompt = line['input'] + first_token + ':'
+    first_token = line["output"].split(":")[0]
+    prompt = line["input"] + first_token + ":"
 
-    input_ids = tokenizer.encode(prompt, return_tensors='pt').cuda()[:, : 8192 - 1024]
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").cuda()[:, : 8192 - 1024]
     output = vardecoder_model.generate(
-        input_ids=input_ids, max_new_tokens=1024, num_beams=4, num_return_sequences=1, do_sample=False,
-        early_stopping=False, pad_token_id=0, eos_token_id=0
+        input_ids=input_ids,
+        max_new_tokens=1024,
+        num_beams=4,
+        num_return_sequences=1,
+        do_sample=False,
+        early_stopping=False,
+        pad_token_id=0,
+        eos_token_id=0,
     )[0]
-    output = tokenizer.decode(output[input_ids.size(1): ], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    output = tokenizer.decode(
+        output[input_ids.size(1) :],
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True,
+    )
 
-    output = first_token + ':' + output
+    output = first_token + ":" + output
     return output
 
-demo = gr.Interface(fn=infer, inputs=gr.Textbox(lines=10, placeholder=example), outputs=gr.Text())
+
+demo = gr.Interface(
+    fn=infer,
+    inputs=[
+        gr.Text(label="First Token", placeholder="a1"),
+        gr.Textbox(lines=10, placeholder=example),
+    ],
+    outputs=gr.Text(label="Var Decoder Output"),
+)
 demo.launch()