Update app.py

app.py

@@ -1,119 +1,112 @@
 import gradio as gr
-import re
 import requests
 import json
 import os
-from screenshot import BG_COMP, BOX_COMP, GENERATION_VAR, PROMPT_VAR, main
-from pathlib import Path
-
-title = "BLOOM"
-description = """Gradio Demo for BLOOM. To use it, simply add your text, or click one of the examples to load them.
-Tips:
-- Do NOT talk to BLOOM as an entity, it's not a chatbot but a webpage/blog/article completion model.
-- For the best results: MIMIC a few sentences of a webpage similar to the content you want to generate.
-Start a paragraph as if YOU were writing a blog, webpage, math post, coding article and BLOOM will generate a coherent follow-up. Longer prompts usually give more interesting results.
-- Content: Please see our [content disclaimer](https://huggingface.co/spaces/bigscience/bloom-book) before using the model, as it may sometimes behave in unexpected ways.
-
-Options:
-- sampling: imaginative completions (may be not super accurate e.g. math/history)
-- greedy: accurate completions (may be more boring or have repetitions)
-"""
+from screenshot import (
+    before_prompt,
+    prompt_to_generation,
+    after_generation,
+    js_save,
+    js_load_script,
+)
+from spaces_info import description, examples, initial_prompt_value
 
 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, "Sample", False, "Sample 1"],
-    ['A poem about the beauty of science by Alfred Edgar Brittle\nTitle: The Magic Craft\nIn the old times', 50, "Sample", False,  "Sample 1"],
-    ['استخراج العدد العاملي في لغة بايثون:', 30, "Greedy", False, "Sample 1"],
-    ["Pour déguster un ortolan, il faut tout d'abord", 32, "Sample", False, "Sample 1"],
-    ['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 -', 16, "Sample", False, "Sample 1"],
-    ['Estos ejemplos quitan vocales de las palabras\nEjemplos:\nhola - hl\nmanzana - mnzn\npapas - pps\nalacran - lcrn\npapa -', 16, "Sample",False,  "Sample 1"],
-    ["Question: If I put cheese into the fridge, will it melt?\nAnswer:", 32, "Sample", False, "Sample 1"],
-    ["Math exercise - answers:\n34+10=44\n54+20=", 16, "Greedy", False,  "Sample 1"],
-    ["Question: Where does the Greek Goddess Persephone spend half of the year when she is not with her mother?\nAnswer:", 24, "Greedy", False, "Sample 1"],
-    ["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:", 24, "Greedy", False, "Sample 1"],
-]
 
 def query(payload):
     print(payload)
     response = requests.request("POST", API_URL, json=payload)
     print(response)
     return json.loads(response.content.decode("utf-8"))
-    
-def inference(input_sentence, max_length, sample_or_greedy, raw_text=False, seed=42):
+
+def inference(input_sentence, max_length, sample_or_greedy, seed=42):
     if sample_or_greedy == "Sample":
-        parameters = {"max_new_tokens": max_length,
-                      "top_p": 0.9,
-                      "do_sample": True,
-                      "seed": seed,
-                      "early_stopping": False,
-                      "length_penalty": 0.0,
-                      "eos_token_id": None}
+        parameters = {
+            "max_new_tokens": max_length,
+            "top_p": 0.9,
+            "do_sample": True,
+            "seed": seed,
+            "early_stopping": False,
+            "length_penalty": 0.0,
+            "eos_token_id": None,
+        }
     else:
-        parameters = {"max_new_tokens": max_length,
-                      "do_sample": False,
-                      "seed": seed,
-                      "early_stopping": False,
-                      "length_penalty": 0.0,
-                      "eos_token_id": None}
+        parameters = {
+            "max_new_tokens": max_length,
+            "do_sample": False,
+            "seed": seed,
+            "early_stopping": False,
+            "length_penalty": 0.0,
+            "eos_token_id": None,
+        }
 
-    payload = {"inputs": input_sentence,
-               "parameters": parameters}
+    payload = {"inputs": input_sentence, "parameters": parameters}
 
-    data = query(
-        payload
-    )
-    
-    if raw_text:
-        return None, data[0]['generated_text']
-    
-    width, height = 3246, 3246
-    assets_path = "assets"
-    font_mapping = {
-        "latin characters (faster)": "DejaVuSans.ttf",
-        "complete alphabet (slower)":"GoNotoCurrent.ttf"
-    }
-    working_dir = Path(__file__).parent.resolve()
-    font_path = str(working_dir / font_mapping["complete alphabet (slower)"])
-    img_save_path = str(working_dir / "output.jpeg")
-    colors = {
-        BG_COMP: "#000000",
-        PROMPT_VAR: "#FFFFFF",
-        GENERATION_VAR: "#FF57A0",
-        BOX_COMP: "#120F25",
-    }
-    
-    new_string = data[0]['generated_text'].split(input_sentence, 1)[1]
+    #if API_URL is not None:
+    data = query(payload)
+    #else:  # To remove when plugged in a real API
+    #    input_sentence, generation = payload["inputs"].split("\n", 1)
+    #    data = [{"generated_text": input_sentence + generation}]
 
-    _, img = main(
-        input_sentence,
-        new_string,
-        width,
-        height,
-        assets_path=assets_path,
-        font_path=font_path,
-        colors=colors,
-        frame_to_box_margin=200,
-        text_to_text_box_margin=50,
-        init_font_size=150,
-        right_align=False,
+    generation = data[0]["generated_text"].split(input_sentence, 1)[1]
+    return (
+        before_prompt
+        + input_sentence
+        + prompt_to_generation
+        + generation
+        + after_generation,
+        data[0]["generated_text"],
     )
-    return img, data[0]['generated_text']
 
 
-gr.Interface(
-    inference, 
-    [
-        gr.inputs.Textbox(label="Input"),
-        gr.inputs.Slider(1, 64, default=32, step=1, label="Tokens to generate"),
-        gr.inputs.Radio(["Sample", "Greedy"], label="Sample or greedy", default="Sample"),
-        gr.Checkbox(label="Just output raw text"),
-        gr.inputs.Radio(["Sample 1", "Sample 2", "Sample 3", "Sample 4", "Sample 5"], default="Sample 1", label="Sample other generations (only work in 'Sample' mode", type="index"),
-    ], 
-    ["image", "text"],
-    examples=examples,
-    # article=article,
-    cache_examples=False,
-    title=title,
-    description=description
-).launch()
+if __name__ == "__main__":
+    demo = gr.Blocks()
+    with demo:
+        with gr.Row():
+            gr.Markdown(value=description)
+        with gr.Row():
+            with gr.Column():
+                text = gr.Textbox(
+                    label="Input",
+                    value=" ",  # should be set to " " when plugged into a real API
+                )
+                tokens = gr.Slider(1, 64, value=32, step=1, label="Tokens to generate")
+                sampling = gr.Radio(
+                    ["Sample", "Greedy"], label="Sample or greedy", value="Sample"
+                )
+                sampling2 = gr.Radio(
+                    ["Sample 1", "Sample 2", "Sample 3", "Sample 4", "Sample 5"],
+                    value="Sample 1",
+                    label="Sample other generations (only work in 'Sample' mode)",
+                    type="index",
+                )
+
+                with gr.Row():
+                    submit = gr.Button("Submit")
+                    load_image = gr.Button("Generate Image")
+            with gr.Column():
+                display_out = gr.HTML(
+                    label="Image"
+                )  # , visible=False) #interactive=True,
+                display_out.set_event_trigger(
+                    "load",
+                    fn=None,
+                    inputs=None,
+                    outputs=None,
+                    no_target=True,
+                    js=js_load_script,
+                )
+                text_out = gr.Textbox(label="Output")
+        with gr.Row():
+           gr.Examples(examples=examples, inputs=[text, tokens, sampling, sampling2])
+
+        submit.click(
+            inference,
+            inputs=[text, tokens, sampling, sampling2],
+            outputs=[display_out, text_out],
+        )
+
+        load_image.click(fn=None, inputs=None, outputs=None, _js=js_save)
+
+    demo.launch()