Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 import spaces
 import gradio as gr
+import flash_attn
 from threading import Thread
 from transformers import (
     AutoModelForCausalLM, 
@@ -11,7 +12,7 @@ from transformers import (
     StoppingCriteriaList
 )
 
-MODEL_ID = "cognitivecomputations/Dolphin3.0-R1-Mistral-24B"
+MODEL_ID = "NousResearch/DeepHermes-3-Llama-3-8B-Preview"
 
 DEFAULT_SYSTEM_PROMPT = """
 
@@ -61,7 +62,8 @@ def initialize_model():
         device_map="cuda",
        # quantization_config=quantization_config,
         torch_dtype=torch.bfloat16,
-        trust_remote_code=True
+        trust_remote_code=True,
+        attn_implementation="flash_attention_2"
     )
     model.to("cuda")
     model.eval()  # set evaluation mode to disable gradients and speed up inference