Update app.py

app.py

@@ -20,49 +20,62 @@ Options:
 
 wip_description = """JAX / Flax Gradio Demo for BLOOM. The 176B BLOOM model running on a TPU v3-256 pod, with 2D model parallelism and custom mesh axes.
 
-Note: for this WIP demo only `sampling` is supported.
+Note: for this WIP demo only **sampling** is supported.
 """
 
-API_URL =  os.getenv("API_URL")
+API_URL = os.getenv("API_URL")
 
 examples = [
-    ['A "whatpu" is a small, furry animal native to Tanzania. An example of a sentence that uses the word whatpu is: We were traveling in Africa and we saw these very cute whatpus. To do a "farduddle" means to jump up and down really fast. An example of a sentence that uses the word farduddle is:', 32, "sampling", False],
-    ['A poem about the beauty of science by Alfred Edgar Brittle\nTitle: The Magic Craft\nIn the old times', 64, "sampling", False],
-    ['استخراج العدد العاملي في لغة بايثون:', 32, "greedy", False],
-    ["Pour déguster un ortolan, il faut tout d'abord", 32, "sampling", False],
-    ['Traduce español de España a español de Argentina\nEl coche es rojo - el auto es rojo\nEl ordenador es nuevo - la computadora es nueva\nel boligrafo es negro -', 32, "sampling", False],
-    ['Estos ejemplos quitan vocales de las palabras\nEjemplos:\nhola - hl\nmanzana - mnzn\npapas - pps\nalacran - lcrn\npapa -', 32, "sampling", False],
-    ["Question: If I put cheese into the fridge, will it melt?\nAnswer:", 32, "sampling", False],
-    ["Math exercise - answers:\n34+10=44\n54+20=", 32, "greedy", False],
-    ["Question: Where does the Greek Goddess Persephone spend half of the year when she is not with her mother?\nAnswer:", 32, "greedy", False],
-    ["spelling test answers.\nWhat are the letters in « language »?\nAnswer: l-a-n-g-u-a-g-e\nWhat are the letters in « Romanian »?\nAnswer:", 32, "greedy", False]
+    [
+        'A "whatpu" is a small, furry animal native to Tanzania. An example of a sentence that uses the word whatpu is: We were traveling in Africa and we saw these very cute whatpus. To do a "farduddle" means to jump up and down really fast. An example of a sentence that uses the word farduddle is:',
+        64, "sampling", False],
+    ['A poem about the beauty of science by Alfred Edgar Brittle\nTitle: The Magic Craft\nIn the old times', 64,
+     "sampling", False],
+    ['استخراج العدد العاملي في لغة بايثون:', 64, "greedy", False],
+    ["Pour déguster un ortolan, il faut tout d'abord", 64, "sampling", False],
+    [
+        'Traduce español de España a español de Argentina\nEl coche es rojo - el auto es rojo\nEl ordenador es nuevo - la computadora es nueva\nel boligrafo es negro -',
+        64, "sampling", False],
+    [
+        'Estos ejemplos quitan vocales de las palabras\nEjemplos:\nhola - hl\nmanzana - mnzn\npapas - pps\nalacran - lcrn\npapa -',
+        64, "sampling", False],
+    ["Question: If I put cheese into the fridge, will it melt?\nAnswer:", 64, "sampling", False],
+    ["Math exercise - answers:\n34+10=44\n54+20=", 64, "greedy", False],
+    [
+        "Question: Where does the Greek Goddess Persephone spend half of the year when she is not with her mother?\nAnswer:",
+        64, "greedy", False],
+    [
+        "spelling test answers.\nWhat are the letters in « language »?\nAnswer: l-a-n-g-u-a-g-e\nWhat are the letters in « Romanian »?\nAnswer:",
+        64, "greedy", False]
 ]
 
+
 def query(payload):
     print(payload)
     response = requests.post(API_URL, json=payload)
     print(response)
     return response.json()
-    
+
+
 def inference(input_sentence, max_length, sample_or_greedy, raw_text=False):
     payload = {
-    "prompt": input_sentence, 
-    "do_sample": True,
-    #"max_new_tokens": max_length
+        "prompt": input_sentence,
+        "do_sample": True,
+        # "max_new_tokens": max_length
     }
 
     data = query(
         payload
     )
-    
+
     if raw_text:
-        return None, data['generated_text'][0]
-    
+        return None, data['generated_text']
+
     width, height = 3326, 3326
     assets_path = "assets"
     font_mapping = {
         "latin characters (faster)": "DejaVuSans.ttf",
-        "complete alphabet (slower)":"GoNotoCurrent.ttf"
+        "complete alphabet (slower)": "GoNotoCurrent.ttf"
     }
     working_dir = Path(__file__).parent.resolve()
     font_path = str(working_dir / font_mapping["complete alphabet (slower)"])
@@ -73,8 +86,8 @@ def inference(input_sentence, max_length, sample_or_greedy, raw_text=False):
         GENERATION_VAR: "#FF57A0",
         BOX_COMP: "#120F25",
     }
-    
-    new_string = data['generated_text'][0]
+
+    new_string = data['generated_text']
 
     _, img = main(
         input_sentence,
@@ -93,13 +106,13 @@ def inference(input_sentence, max_length, sample_or_greedy, raw_text=False):
 
 
 gr.Interface(
-    inference, 
+    inference,
     [
         gr.inputs.Textbox(label="Input"),
         gr.inputs.Radio([64], default=64, label="Tokens to generate"),
         gr.inputs.Radio(["sampling"], label="Sample or greedy", default="sampling"),
         gr.Checkbox(label="Just output raw text", value=True),
-    ], 
+    ],
     ["image", "text"],
     examples=examples,
     # article=article,