Update app.py

app.py

@@ -1,10 +1,15 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import subprocess
+import sys
+
+# Force install the specific transformers version from the GitHub PR
+subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "git+https://github.com/Muennighoff/transformers.git@olmoe"])
 
-# Install flash attention
-subprocess.run('pip install --upgrade --force-reinstall --no-deps --no-build-isolation transformers torch flash-attn  ', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+# Install flash-attn
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Load model and tokenizer
 model_name = "allenai/OLMoE-1B-7B-0924-Instruct"