fixes

app.py

@@ -1,8 +1,7 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import spaces
-from PIL import Image
-import io
+import re
 
 
 models = {
@@ -20,9 +19,11 @@ def run_example(html_content, model_id="jinaai/reader-lm-0.5b"):
     tokenizer = tokenizers[model_id]
     messages = [{"role": "user", "content": html_content}]
     input_text=tokenizer.apply_chat_template(messages, tokenize=False)
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to("cuda")
     outputs = model.generate(inputs, max_new_tokens=1024, temperature=0, do_sample=False, repetition_penalty=1.08)
-    return tokenizer.decode(outputs[0])
+    pattern = r"<\|im_start\|>assistant(.*?)<\|im_end\|>"
+    assistant_response = re.findall(pattern, tokenizer.decode(outputs[0]), re.DOTALL)
+    return assistant_response[0]
 
 
 css = """
@@ -40,7 +41,7 @@ with gr.Blocks(css=css) as demo:
     with gr.Tab(label="Main"):
         with gr.Row():
             with gr.Column():
-                model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value="maxiw/Florence-2-ScreenQA-base")
+                model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value="jinaai/reader-lm-0.5b")
                 html_content = gr.Textbox(label="HTML")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():