Fix dependencies and styles
diff --git a/fastchat/llm_judge/README.md b/fastchat/llm_judge/README.md
index 1ab0fdc9d2..b9f06820bf 100644
--- a/fastchat/llm_judge/README.md
+++ b/fastchat/llm_judge/README.md
@@ -16,7 +16,7 @@ To automate the evaluation process, we prompt strong LLMs like GPT-4 to act as j
 ```
 git clone https://github.com/lm-sys/FastChat.git
 cd FastChat
-pip install -e ".[llm_judge]"
+pip install -e ".[model_worker,llm_judge]"
 ```
 
 ## Review Pre-Generated Model Answers and Judgments
diff --git a/fastchat/model/model_falcon.py b/fastchat/model/model_falcon.py
index 740f6ef5bb..20afc4f0fb 100644
--- a/fastchat/model/model_falcon.py
+++ b/fastchat/model/model_falcon.py
@@ -8,8 +8,6 @@
 
 from fastchat.utils import is_partial_stop
 
-transformers.logging.set_verbosity_error()
-
 
 @torch.inference_mode()
 def generate_stream_falcon(
diff --git a/fastchat/train/llama_flash_attn_monkey_patch.py b/fastchat/train/llama_flash_attn_monkey_patch.py
index 89f1913d97..5f9971a32e 100644
--- a/fastchat/train/llama_flash_attn_monkey_patch.py
+++ b/fastchat/train/llama_flash_attn_monkey_patch.py
@@ -1,26 +1,14 @@
 from typing import List, Optional, Tuple
-import logging
+import warnings
 
 import torch
 from torch import nn
-import warnings
 import transformers
 from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
 
-try:
-    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
-    from flash_attn.bert_padding import unpad_input, pad_input
-except Exception:
-    raise ModuleNotFoundError(
-        "Please install FlashAttention first, e.g., with pip install flash-attn --no-build-isolation, Learn more at https://github.com/Dao-AILab/flash-attention#installation-and-features"
-    )
-
-try:
-    from einops import rearrange
-except Exception:
-    raise ModuleNotFoundError(
-        "Please install einops first, e.g., with pip install einops"
-    )
+from einops import rearrange
+from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
+from flash_attn.bert_padding import unpad_input, pad_input
 
 
 # ADAPTED from https://github.com/allenai/open-instruct/blob/main/open_instruct/llama_flash_attn_monkey_patch.py
@@ -135,7 +123,7 @@ def _prepare_decoder_attention_mask(
 def replace_llama_attn_with_flash_attn():
     cuda_major, cuda_minor = torch.cuda.get_device_capability()
     if cuda_major < 8:
-        logging.warning(
+        warnings.warn(
             "Flash attention is only supported on A100 or H100 GPU during training due to head dim > 64 backward."
             "ref: https://github.com/HazyResearch/flash-attention/issues/190#issuecomment-1523359593"
         )
diff --git a/pyproject.toml b/pyproject.toml
index 86fb5cd315..86ff61b85a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,15 +13,15 @@ classifiers = [
     "License :: OSI Approved :: Apache Software License",
 ]
 dependencies = [
-    "einops", "fastapi", "httpx", "markdown2[all]", "nh3", "numpy",
-    "peft", "prompt_toolkit>=3.0.0", "pydantic<2,>=1", "requests", "rich>=10.0.0",
+    "fastapi", "httpx", "markdown2[all]", "nh3", "numpy",
+    "prompt_toolkit>=3.0.0", "pydantic<2,>=1", "requests", "rich>=10.0.0",
     "shortuuid", "tiktoken", "uvicorn", "wandb"
 ]
 
 [project.optional-dependencies]
-model_worker = ["accelerate>=0.21", "torch", "transformers>=4.31.0"]
+model_worker = ["accelerate>=0.21", "peft", "sentencepiece", "torch", "transformers>=4.31.0"]
 webui = ["gradio"]
-train = ["flash-attn>=2.0"]
+train = ["einops", "flash-attn>=2.0"]
 llm_judge = ["openai", "anthropic>=0.3", "ray"]
 dev = ["black==23.3.0", "pylint==2.8.2"]
 
