Update app.py

app.py

@@ -2,112 +2,48 @@ import gradio as gr
 import requests
 import json
 import os
-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")
 HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 
+# Global variable to store the generated text
+generated_text = ""
 
 def query(payload):
-    print(payload)
     response = requests.request("POST", API_URL, json=payload, headers={"Authorization": f"Bearer {HF_API_TOKEN}"})
-    print(response)
     return json.loads(response.content.decode("utf-8"))
 
-
-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,
-        }
-    else:
-        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,"options" : {"use_cache": False} }
-
+def generate_and_append_text(max_length):
+    global generated_text
+    parameters = {
+        "max_new_tokens": max_length,
+        "top_p": 0.9,
+        "do_sample": True,
+        "seed": 42,
+        "early_stopping": False,
+        "length_penalty": 0.0,
+        "eos_token_id": None,
+    }
+    payload = {"inputs": generated_text, "parameters": parameters, "options": {"use_cache": False}}
     data = query(payload)
-
     if "error" in data:
-        return (None, None, f"<span style='color:red'>ERROR: {data['error']} </span>")
-
-    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 f"<span style='color:red'>ERROR: {data['error']} </span>"
+    new_text = data[0]["generated_text"].replace(generated_text, "").strip()
+    generated_text += " " + new_text
+    return generated_text
 
 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():
-                text_error = gr.Markdown(label="Log information")
-                text_out = gr.Textbox(label="Output")
-                display_out = gr.HTML(label="Image")
-                display_out.set_event_trigger(
-                    "load",
-                    fn=None,
-                    inputs=None,
-                    outputs=None,
-                    no_target=True,
-                    js=js_load_script,
-                )
-        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, text_error],
+            generate_button = gr.Button("Generate Text")
+            tokens = gr.Slider(1, 64, value=32, step=1, label="Tokens to generate")
+            text_out = gr.Textbox(label="Generated Text")
+
+        generate_button.click(
+            generate_and_append_text,
+            inputs=tokens,
+            outputs=text_out
         )
 
-        load_image.click(fn=None, inputs=None, outputs=None, _js=js_save)
-
     demo.launch()