Remove think block from response

app.py

@@ -18,6 +18,11 @@ def generate_response(prompt):
     
     # Decode the generated text
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Remove the think block
+    ix = generated_text.find("</think>")
+    if ix != -1:
+        generated_text = generated_text[ix + len("</think>") :].lstrip()
         
     return generated_text
 
@@ -42,4 +47,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()