a

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from threading import Thread
 
 from marker.convert import convert_single_pdf
 from marker.output import markdown_exists, save_markdown, get_markdown_filepath
@@ -33,6 +34,7 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
 
 # Define prompts
 SYSTEM_PROMPT_TEMPLATE = """You are an expert reviewer for AI conferences. You follow best practices and review papers according to the reviewer guidelines.
@@ -172,7 +174,7 @@ def process_file(file):
     return paper_text
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU()
 def generate(paper_text, review_template):
     # messages = [
     #     {"role": "system", "content": "You are a pirate."},
@@ -185,19 +187,26 @@ def generate(paper_text, review_template):
         return_tensors='pt'
     ).to(model.device)
     print(f"input_ids shape: {input_ids.shape}")
-    generated_ids = model.generate(
-        input_ids=input_ids,
-        max_new_tokens=4096,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.9
-    )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(input_ids, generated_ids)
-    ]
+    generation_kwargs = dict(input_ids=input_ids, streamer=streamer, max_new_tokens=4096, do_sample=True, temperature=0.6, top_p=0.9)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text
+    # generated_ids = model.generate(
+    #     input_ids=input_ids,
+    #     max_new_tokens=4096,
+    #     do_sample=True,
+    #     temperature=0.6,
+    #     top_p=0.9
+    # )
+    # generated_ids = [
+    #     output_ids[len(input_ids):] for input_ids, output_ids in zip(input_ids, generated_ids)
+    # ]
     
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return response
+    # response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    # return response
     # return "Success"
 
 
@@ -249,13 +258,13 @@ with gr.Blocks(theme=theme) as demo:
     review_field = gr.Markdown("\n\n\n\n\n", label="Review")
     generate_button.click(fn=lambda: gr.update(interactive=False), inputs=None, outputs=generate_button).then(generate, [paper_text_field, review_template_field], review_field).then(fn=lambda: gr.update(interactive=True), inputs=None, outputs=generate_button)
     
-    gr.Examples([
-        ["examples/2105.04505v1.pdf", REVIEW_FIELDS]
-    ],
-    inputs=[paper_text_field, review_template_field],
-    outputs=[review_field],
-    fn=generate,
-    cache_examples=True)
+    # gr.Examples([
+    #     ["examples/2105.04505v1.pdf", REVIEW_FIELDS]
+    # ],
+    # inputs=[paper_text_field, review_template_field],
+    # outputs=[review_field],
+    # fn=generate,
+    # cache_examples=True)
 
     demo.title = "Paper Review Generator"