visualize beam in output

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 def translate(input_txt,prefix,suffix,num_beams,topk,beta):
     dictToSend = {'input':input_txt,'prefix':prefix,'suffix':suffix,'num_beams':num_beams,'alpha': (1-beta),'beta':beta,'topk':topk}
     res = requests.post('https://d7bf-141-3-25-29.ngrok-free.app/generate', json=dictToSend)
-    return res.text
+    return res.text,res.beam_img
 
 prefix= "<s>[INST] <<SYS>>\nYou are a professional translator from English to German. You translate the sentences to only Formal German even if the English sentence is Informal.\n<</SYS>>\nEnglish: "
 #self.prefix = "Translate from English to German:\nEnglish: "
@@ -21,8 +21,8 @@ example_informal = ["Where are you going today?",prefix,suffix,5,5,1]
 
 iface = gr.Interface(
         fn=translate,
-        inputs=[gr.Textbox(lines=1, placeholder="Enter your English Sentences that you want to translate", label="English Sentence"), gr.Textbox(lines=5, placeholder=prefix,label="Prefix for LLM Prompt - Replace formal and informal here to control Formality Levle"),gr.Textbox(lines=2, placeholder=suffix, label="DONT CHANGE!!"),gr.Slider(label='Select beam size',value=5,minimum=1,maximum=5,step=1),gr.Slider(label="TopK for rescoring",value=5,minimum=5,maximum=10,step=1),gr.Slider(label="LLM weight when rescoring, 1 for complete LLM scoring, 0 for NLLB Decoding without LLM",value=1,minimum=0,maximum=1,step=0.1)],
-        outputs=gr.Textbox(lines=1,placeholder="Enter your inputs and click submit!",label="LLM Guided NLLB German Translation"),
+        inputs=[gr.Textbox(lines=1, placeholder="Enter your English Sentences that you want to translate", label="English Sentence"), gr.Textbox(lines=5, placeholder=prefix,label="Prefix for LLM Prompt - Replace formal and informal here to control Formality Level"),gr.Textbox(lines=2, placeholder=suffix, label="DONT CHANGE!!"),gr.Slider(label='Select beam size',value=5,minimum=1,maximum=5,step=1),gr.Slider(label="TopK for rescoring",value=5,minimum=5,maximum=10,step=1),gr.Slider(label="LLM weight when rescoring, 1 for complete LLM scoring, 0 for NLLB Decoding without LLM",value=1,minimum=0,maximum=1,step=0.1)],
+        outputs=[gr.Textbox(lines=1,placeholder="Enter your inputs and click submit!",label="LLM Guided NLLB German Translation"),gr.Image(label="Beam Visualization")],
         examples=[example_formal,example_informal],
         title="Ensembling NLLB 3.3B with LLama2 13B by Guided Decoding",
         )