Ruff check and format

Signed-off-by: Amit Raj <quic_amitraj@quicinc.com>

Author: quic-amitraj

QEfficient/transformers/embeddings/embedding_utils.py

@@ -57,6 +57,7 @@ def forward(
         output = self.base_model(input_ids, attention_mask, **kwargs)
         return self.pooling_fn(output[0], attention_mask)
 
+
 def validate_user_pooling_function(user_function):
     if not callable(user_function):
         raise TypeError("Provided pooling function is not callable.")
@@ -65,4 +66,4 @@ def validate_user_pooling_function(user_function):
     required_args = {"last_hidden_states", "attention_mask"}
     if not required_args.issubset(sig.parameters.keys()):
         raise ValueError(f"Pooling function must accept arguments: {required_args}")
-    return user_function
+    return user_function

QEfficient/transformers/models/pytorch_transforms.py

@@ -502,7 +502,9 @@ def apply(cls, model: nn.Module, **kwargs) -> Tuple[nn.Module, bool]:
         transformed = False
         if kwargs.get("pooling") is not None:
             pooling = kwargs["pooling"]
-            pooling_method = POOLING_MAP[pooling] if isinstance(pooling,str) else validate_user_pooling_function(pooling)
+            pooling_method = (
+                POOLING_MAP[pooling] if isinstance(pooling, str) else validate_user_pooling_function(pooling)
+            )
             model = PooledModel(model, pooling_method)
             warnings.warn(f"Pooling method {pooling.__name__} is applied to the model.")
         return model, transformed

examples/embedding_model.py

@@ -13,11 +13,13 @@
 
 from QEfficient import QEFFAutoModel as AutoModel
 
+
 def max_pooling(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
     input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_states.size()).float()
     last_hidden_states[input_mask_expanded == 0] = -1e9
     return torch.max(last_hidden_states, 1)[0]
 
+
 # Sentences we want sentence embeddings for
 sentences = "This is an example sentence"
 
@@ -28,7 +30,7 @@ def max_pooling(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor)
 qeff_model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2", pooling=max_pooling)
 
 # Here seq_len can be list seq_len or single int
-qeff_model.compile(num_cores=16, seq_len=[32,64])
+qeff_model.compile(num_cores=16, seq_len=[32, 64])
 
 # Tokenize sentences
 encoded_input = tokenizer(sentences, return_tensors="pt")