fix generation parsing

app.py

@@ -19,7 +19,7 @@ model = Idefics2ForConditionalGeneration.from_pretrained(
 
 @spaces.GPU(duration=180)
 def model_inference(
-    image, text, decoding_strategy, temperature, 
+    image, text, decoding_strategy, temperature,
     max_new_tokens, repetition_penalty, top_p
 ):
     if text == "" and not image:
@@ -36,16 +36,16 @@ def model_inference(
                 ]
             }
         ]
-    
-    
+
+
     prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
     inputs = processor(text=prompt, images=[image], return_tensors="pt")
     inputs = {k: v.to("cuda") for k, v in inputs.items()}
-    
+
     generation_args = {
         "max_new_tokens": max_new_tokens,
         "repetition_penalty": repetition_penalty,
-        
+
     }
 
     assert decoding_strategy in [
@@ -59,20 +59,15 @@ def model_inference(
         generation_args["do_sample"] = True
         generation_args["top_p"] = top_p
 
-    
+
     generation_args.update(inputs)
 
     # Generate
     generated_ids = model.generate(**generation_args)
-    
-    generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)
-    print(generated_texts)
-    pattern = r"Assistant: (.*)"
 
-    # Use regular expression to find the desired part
-    result = re.search(pattern, generated_texts[0]).group(1)
-
-    return result[:-1]
+    generated_texts = processor.batch_decode(generated_ids[:, generation_args["input_ids"].size(1):], skip_special_tokens=True)
+    print("INPUT:", prompt, "|OUTPUT:", generated_texts)
+    return generated_texts[0]
 
 
 with gr.Blocks(fill_height=True) as demo:
@@ -87,7 +82,7 @@ with gr.Blocks(fill_height=True) as demo:
         query_input = gr.Textbox(label="Prompt")
         submit_btn = gr.Button("Submit")
         output = gr.Textbox(label="Output")
-        
+
     with gr.Accordion(label="Example Inputs and Advanced Generation Parameters"):
         examples=[["./example_images/docvqa_example.png", "How many items are sold?", "Greedy", 0.4, 512, 1.2, 0.8],
                     ["./example_images/example_images_travel_tips.jpg", "I want to go somewhere similar to the one in the photo. Give me destinations and travel tips.", "Greedy", 0.4, 512, 1.2, 0.8],
@@ -95,7 +90,7 @@ with gr.Blocks(fill_height=True) as demo:
                     ["./example_images/dummy_pdf.png", "How much percent is the order status?", "Greedy", 0.4, 512, 1.2, 0.8],
                     ["./example_images/art_critic.png", "As an art critic AI assistant, could you describe this painting in details and make a thorough critic?.", "Greedy", 0.4, 512, 1.2, 0.8],
                     ["./example_images/s2w_example.png", "What is this UI about?", "Greedy", 0.4, 512, 1.2, 0.8]]
-        
+
         # Hyper-parameters for generation
         max_new_tokens = gr.Slider(
               minimum=8,
@@ -151,7 +146,7 @@ with gr.Blocks(fill_height=True) as demo:
               inputs=decoding_strategy,
               outputs=temperature,
           )
-        
+
         decoding_strategy.change(
               fn=lambda selection: gr.Slider(
                   visible=(
@@ -168,13 +163,13 @@ with gr.Blocks(fill_height=True) as demo:
           )
         gr.Examples(
                         examples = examples,
-                        inputs=[image_input, query_input, decoding_strategy, temperature, 
+                        inputs=[image_input, query_input, decoding_strategy, temperature,
                                                               max_new_tokens, repetition_penalty, top_p],
                         outputs=output,
                         fn=model_inference
                     )
-          
-        submit_btn.click(model_inference, inputs = [image_input, query_input, decoding_strategy, temperature, 
+
+        submit_btn.click(model_inference, inputs = [image_input, query_input, decoding_strategy, temperature,
                                                       max_new_tokens, repetition_penalty, top_p], outputs=output)