Update app.py

app.py

@@ -6,7 +6,7 @@ import sys
 
 # Install required packages
 subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "--force-reinstall", "--no-deps", "einops", "accelerate", "torch", "git+https://github.com/Muennighoff/transformers.git@olmoe"])
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+#subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 from transformers import OlmoeForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
@@ -19,10 +19,10 @@ try:
     model = OlmoeForCausalLM.from_pretrained(
         model_name, 
         trust_remote_code=True, 
-        torch_dtype=torch.float16,  # Using float16 for lower precision
+        torch_dtype=torch.bfloat16,  # Using float16 for lower precision
         low_cpu_mem_usage=True,
         device_map="auto",
-        _attn_implementation="flash_attention_2"  # Enable Flash Attention 2
+        #_attn_implementation="flash_attention_2"  # Enable Flash Attention 2
     ).to(DEVICE)
     model.gradient_checkpointing_enable()  # Enable gradient checkpointing
     tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -88,7 +88,7 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (Now with Flash Attention 2!)")
+    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (CPU experiment)")
     chatbot = gr.Chatbot(elem_id="output")
     msg = gr.Textbox(label="Meow")
     with gr.Row():