ruff

server/lorax_server/models/custom_modeling/flash_cohere_modeling.py

@@ -18,27 +18,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List, Optional, Tuple
+
 import dropout_layer_norm
 import rotary_emb
 import torch
 import torch.distributed
 from torch import nn
 from transformers.activations import ACT2FN
-from typing import Optional, List, Tuple
 
 from lorax_server.adapters.weights import AdapterBatchData
+from lorax_server.utils import flash_attn, paged_attn
 from lorax_server.utils.layers import (
+    FastLayerNorm,
+    MultiAdapterHead,
+    PositionRotaryEmbedding,
     TensorParallelAdapterRowLinear,
-    TensorParallelMultiAdapterLinear,
-    TensorParallelRowLinear,
     TensorParallelColumnLinear,
     TensorParallelEmbedding,
-    PositionRotaryEmbedding,
+    TensorParallelHead,
+    TensorParallelMultiAdapterLinear,
+    TensorParallelRowLinear,
     get_linear,
-    FastLayerNorm,
 )
-from lorax_server.utils import flash_attn, paged_attn
-from lorax_server.utils.layers import MultiAdapterHead, TensorParallelHead
 from lorax_server.utils.lora import DOWN_PROJ, GATE_PROJ, K_PROJ, LM_HEAD, O_PROJ, Q_PROJ, UP_PROJ, V_PROJ
 
 

server/lorax_server/models/flash_cohere.py

@@ -1,18 +1,18 @@
+from typing import Dict, List, Optional, Tuple
+
 import torch
 import torch.distributed
-
 from opentelemetry import trace
-from typing import Dict, List, Optional, Tuple
-from transformers import AutoTokenizer, AutoConfig
+from transformers import AutoConfig, AutoTokenizer
 
 from lorax_server.models import FlashCausalLM
 from lorax_server.models.custom_modeling.flash_cohere_modeling import (
     FlashCohereForCausalLM,
 )
 from lorax_server.utils import (
+    Weights,
     initialize_torch_distributed,
     weight_files,
-    Weights,
 )
 from lorax_server.utils.lora import DOWN_PROJ, GATE_PROJ, K_PROJ, O_PROJ, Q_PROJ, UP_PROJ, V_PROJ