Add break after end of sequence token

app.py

@@ -62,7 +62,7 @@ with gr.Blocks() as demo:
             yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
             return
 
-        current_model.to("cuda")
+        current_model = current_model.half().to("cuda")
 
         prompt = format_prompt(messages)
         inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
@@ -79,10 +79,18 @@ with gr.Blocks() as demo:
             output_scores=False
         ).sequences[0][inputs['input_ids'].shape[-1]:]:  # skip input tokens
             output_ids.append(token_id.item())
-            decoded = current_tokenizer.decode(output_ids, skip_special_tokens=True)
+            decoded = current_tokenizer.decode(output_ids, skip_special_tokens=False)
+            if output_ids[-1] == current_tokenizer.eos_token_id:
+                current_model.to("cpu")
+                torch.cuda.empty_cache()
+                return
             messages[-1]["content"] = decoded
             yield messages
 
+        current_model.to("cpu")
+        torch.cuda.empty_cache()
+        return
+
     with gr.Row():
         model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
         model_status = gr.Textbox(label="Model Status", interactive=False)