Update app.py

app.py

@@ -20,7 +20,7 @@ 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
@@ -88,7 +88,7 @@ with gr.Blocks(css=css) as demo:
     msg = gr.Textbox(label="Meow")
     with gr.Row():
         temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
-        max_new_tokens = gr.Slider(minimum=50, maximum=4000, value=2000, step=50, label="Max New Tokens")
+        max_new_tokens = gr.Slider(minimum=50, maximum=8000, value=2000, step=50, label="Max New Tokens")
     clear = gr.Button("Clear")
 
     def user(user_message, history):