Update app.py

app.py

@@ -5,10 +5,11 @@ import subprocess
 import sys
 
 # Install required packages
-subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "--force-reinstall", "einops", "accelerate", "git+https://github.com/Muennighoff/transformers.git@olmoe"])
+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)
 
-from transformers import OlmoeForCausalLM, AutoTokenizer
+from transformers import OlmoeForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 
 model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
 
@@ -51,7 +52,7 @@ 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 = gr.TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+        streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
         generation_kwargs = dict(
             inputs=inputs,
             max_new_tokens=max_new_tokens,
@@ -61,7 +62,7 @@ def generate_response(message, history, temperature, max_new_tokens):
             streamer=streamer
         )
 
-        thread = torch.multiprocessing.Process(target=model.generate, kwargs=generation_kwargs)
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
         generated_text = ""