Update app.py

app.py

@@ -3,11 +3,9 @@ import spaces
 import torch
 import subprocess
 import sys
-from threading import Thread
-from transformers import TextIteratorStreamer
 
 # Install required packages
-subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "--force-reinstall", "--no-deps", "einops", "accelerate", "transformers", "sentencepiece", "torch"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "--force-reinstall", "einops", "accelerate", "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)
 
 from transformers import OlmoeForCausalLM, AutoTokenizer
@@ -20,7 +18,7 @@ try:
     model = OlmoeForCausalLM.from_pretrained(
         model_name, 
         trust_remote_code=True, 
-        torch_dtype=torch.bfloat16,  # Using float16 for lower precision
+        torch_dtype=torch.float16,  # Using float16 for lower precision
         low_cpu_mem_usage=True,
         device_map="auto",
         _attn_implementation="flash_attention_2"  # Enable Flash Attention 2
@@ -53,26 +51,32 @@ def generate_response(message, history, temperature, max_new_tokens):
     inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
     
     try:
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        streamer = gr.TextIteratorStreamer(tokenizer, skip_special_tokens=True)
         generation_kwargs = dict(
             inputs=inputs,
-            streamer=streamer,
             max_new_tokens=max_new_tokens,
             do_sample=True,
             temperature=temperature,
             eos_token_id=tokenizer.eos_token_id,
+            streamer=streamer
         )
 
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread = torch.multiprocessing.Process(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
-        partial_message = ""
+        generated_text = ""
         for new_text in streamer:
-            partial_message += new_text
-            yield partial_message.strip()
-
+            generated_text += new_text
+            yield generated_text.strip()
+
+        thread.join()
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            yield "GPU memory exceeded. Try reducing the max tokens or using a smaller model."
+        else:
+            yield f"An error occurred: {str(e)}"
     except Exception as e:
-        yield f"An error occurred: {str(e)}"
+        yield f"An unexpected error occurred: {str(e)}"
 
 css = """
   #output {
@@ -83,12 +87,12 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (Now with Flash Attention 2 and Streaming!)")
+    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (Now with Flash Attention 2!)")
     chatbot = gr.Chatbot(elem_id="output")
     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=8000, value=2000, step=50, label="Max New Tokens")
+        max_new_tokens = gr.Slider(minimum=50, maximum=4000, value=2000, step=50, label="Max New Tokens")
     clear = gr.Button("Clear")
 
     def user(user_message, history):