feat: flexible transformers config

- allow passing arbitrary kwargs to transformers backend
- merge model_kwargs with kwargs for backward compatibility
- document transformers configuration examples (flash attn, custom bnb)

Co-Authored-By: Warp <agent@warp.dev>

Author: j341nono

README.md

@@ -71,6 +71,37 @@ enc = llmembed.Encoder(
 embeddings = enc.encode("Hello world", pooling="pcoteol")
 ```
 
+## Transformers Backend Configuration
+
+When using the `transformers` backend, you can pass standard Hugging Face `AutoModel` arguments directly to the `Encoder`.
+
+**Example 1: Using Flash Attention 2**
+
+```python
+import torch
+
+encoder = Encoder(
+    model_name="meta-llama/Llama-3.1-8B",
+    backend="transformers",
+    attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16
+)
+```
+
+**Example 2: Custom Quantization Config**
+
+```python
+from transformers import BitsAndBytesConfig
+
+bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
+
+encoder = Encoder(
+    model_name="meta-llama/Llama-3.1-8B",
+    backend="transformers",
+    quantization_config=bnb_config
+)
+```
+
 ## vLLM Backend Configuration
 
 When using the `vllm` backend, you can pass native vLLM configuration arguments directly to the `Encoder`. This allows full control over memory usage, parallelism, and model length.

docs/index.md

@@ -71,6 +71,37 @@ enc = llmembed.Encoder(
 embeddings = enc.encode("Hello world", pooling="pcoteol")
 ```
 
+## Transformers Backend Configuration
+
+When using the `transformers` backend, you can pass standard Hugging Face `AutoModel` arguments directly to the `Encoder`.
+
+**Example 1: Using Flash Attention 2**
+
+```python
+import torch
+
+encoder = Encoder(
+    model_name="meta-llama/Llama-3.1-8B",
+    backend="transformers",
+    attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16
+)
+```
+
+**Example 2: Custom Quantization Config**
+
+```python
+from transformers import BitsAndBytesConfig
+
+bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
+
+encoder = Encoder(
+    model_name="meta-llama/Llama-3.1-8B",
+    backend="transformers",
+    quantization_config=bnb_config
+)
+```
+
 ## vLLM Backend Configuration
 
 When using the `vllm` backend, you can pass native vLLM configuration arguments directly to the `Encoder`. This allows full control over memory usage, parallelism, and model length.

src/llmembed/backends/transformers_backend.py

@@ -40,13 +40,16 @@ def __init__(
         self.model = None
         self.tokenizer = None
 
-        # Extract model_kwargs if present
-        model_kwargs = kwargs.pop("model_kwargs", {})
-        self._load_model(model_kwargs)
+        # Merge model_kwargs into kwargs if present for backward compatibility or explicit usage
+        if "model_kwargs" in kwargs:
+            kwargs.update(kwargs.pop("model_kwargs"))
+            
+        self._load_model(kwargs)
 
-    def _load_model(self, model_kwargs: "Dict[str, Any]") -> None:
+    def _load_model(self, load_kws: "Dict[str, Any]") -> None:
         quantization_config = None
-        load_kws = model_kwargs.copy()
+        # Make a copy to avoid mutating the original kwargs if used elsewhere
+        load_kws = load_kws.copy()
 
         if self.quantization:
             if self.quantization == "4bit":