chore: enable flash attention 2

README.md

@@ -1,6 +1,6 @@
 ---
 title: Ghost 8B Beta
-emoji: 👻
+emoji: 👻 / 🥸
 colorFrom: indigo
 colorTo: pink
 sdk: gradio

app.py

@@ -1,3 +1,13 @@
+# pylint: skip-file
+
+import subprocess
+
+subprocess.run(
+    f"pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
+
 import os
 from threading import Thread
 from typing import Iterator
@@ -7,6 +17,7 @@ import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
+
 MAX_MAX_NEW_TOKENS = 4096
 DEFAULT_MAX_NEW_TOKENS = 1536
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))
@@ -18,7 +29,7 @@ DESCRIPTION = """\
 
 The languages supported are 🇺🇸 English, 🇫🇷 French, 🇮🇹 Italian, 🇪🇸 Spanish, 🇵🇹 Portuguese, 🇩🇪 German, 🇻🇳 Vietnamese, 🇰🇷 Korean and 🇨🇳 Chinese.
 
-📋 Note: current model version is "disl-0x5-8k" (10 Jul 2024), context length 8k and current status is "moderating / previewing". For detailed information about the model, see [here](https://ghost-x.org/docs/models/ghost-8b-beta/). Try to experience it the way you want!
+📋 Note: current model version is "disl-0x5" (10 Jul 2024), context length 8k (8192 tokens) and current status is "moderating / previewing". For detailed information about the model, see [here](https://ghost-x.org/docs/models/ghost-8b-beta/). Try to experience it the way you want!
 """
 
 
@@ -241,6 +252,8 @@ if torch.cuda.is_available():
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         device_map="auto",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
         trust_remote_code=True,
         token=model_tk,
     )
@@ -363,4 +376,3 @@ with gr.Blocks(fill_height=True, css="style.css") as demo:
 
 if __name__ == "__main__":
     demo.queue(max_size=20).launch(share=True)
-    # demo.launch(share=True)