Feature/mlora

gliclass/model.py

@@ -795,8 +795,15 @@ def pool_outputs(self, encoder_outputs):
             text_embeddings = nn.functional.normalize(text_embeddings, p=2, dim=-1, eps=self.epsilon)
         return text_embeddings
 
-    def encode_text(self, input_ids, attention_mask):
-        outputs = self.encoder_model(input_ids.squeeze(1), attention_mask=attention_mask.squeeze(1))
+    def encode_text(self, input_ids, attention_mask, adapter_ids=None):
+        encoder_kwargs = {}
+        if adapter_ids is not None:
+            encoder_kwargs["adapter_ids"] = adapter_ids
+        outputs = self.encoder_model(
+            input_ids.squeeze(1),
+            attention_mask=attention_mask.squeeze(1),
+            **encoder_kwargs,
+        )
         text_embeddings = self.pool_outputs(outputs)
         return text_embeddings
 
@@ -843,11 +850,12 @@ def forward(
         output_text_embeddings: bool | None = None,
         output_class_embeddings: bool | None = None,
         return_dict: bool | None = None,
+        adapter_ids: list[str] | None = None,
         **kwargs,
     ) -> Tuple | SequenceClassifierOutput:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        text_embeddings = self.encode_text(input_ids, attention_mask)
+        text_embeddings = self.encode_text(input_ids, attention_mask, adapter_ids=adapter_ids)
         class_embeddings = self.encode_classes(class_input_ids, class_attention_mask, labels_mask)
         logits = self.scorer(text_embeddings, class_embeddings) * self.logit_scale.to(class_embeddings.device)
 
@@ -872,7 +880,7 @@ class GLiClassBiEncoderFused(GLiClassBiEncoder):
     def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
         super().__init__(config, from_pretrained)
 
-    def encode_text(self, input_ids, attention_mask, class_embeddings, labels_mask):
+    def encode_text(self, input_ids, attention_mask, class_embeddings, labels_mask, adapter_ids=None):
         embedding_layer = self.encoder_model.get_input_embeddings()
         inputs_embeds = embedding_layer(input_ids)
 
@@ -884,7 +892,14 @@ def encode_text(self, input_ids, attention_mask, class_embeddings, labels_mask):
         selected_class_embeddings = class_embeddings[labels_batch_indices, labels_indices]
 
         inputs_embeds[batch_indices, class_token_indices] = selected_class_embeddings
-        encoder_outputs = self.encoder_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask.squeeze(1))
+        encoder_kwargs = {}
+        if adapter_ids is not None:
+            encoder_kwargs["adapter_ids"] = adapter_ids
+        encoder_outputs = self.encoder_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask.squeeze(1),
+            **encoder_kwargs,
+        )
 
         post_class_embeddings = torch.zeros_like(class_embeddings)
         post_class_embeddings[labels_batch_indices, labels_indices] = encoder_outputs[0][
@@ -903,14 +918,15 @@ def forward(
         output_text_embeddings: bool | None = None,
         output_class_embeddings: bool | None = None,
         return_dict: bool | None = None,
+        adapter_ids: list[str] | None = None,
         **kwargs,
     ) -> Tuple | SequenceClassifierOutput:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         raw_class_embeddings = self.encode_classes(class_input_ids, class_attention_mask, labels_mask)
 
         encoder_outputs, class_embeddings = self.encode_text(
-            input_ids, attention_mask, raw_class_embeddings, labels_mask
+            input_ids, attention_mask, raw_class_embeddings, labels_mask, adapter_ids=adapter_ids
         )
 
         text_embeddings = self.pool_outputs(encoder_outputs)
@@ -1037,5 +1053,7 @@ def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_mult
         return model_embeds
 
     def forward(self, *args, **kwargs):
+        if kwargs.get("adapter_ids") is None:
+            kwargs.pop("adapter_ids", None)
         outputs = self.model(*args, **kwargs)
         return outputs

gliclass/pipeline.py

@@ -249,6 +249,15 @@ def _normalize_classification_types(
         normalized = self._normalize_classification_type(classification_type)
         return [normalized] * num_texts
 
+    def _normalize_adapter_ids(self, adapter_ids: str | List[str] | None, num_texts: int) -> List[str] | None:
+        if adapter_ids is None:
+            return None
+        if isinstance(adapter_ids, str):
+            return [adapter_ids] * num_texts
+        if len(adapter_ids) != num_texts:
+            raise ValueError("Length of adapter_ids list must match number of texts.")
+        return adapter_ids
+
     def _process_labels(
         self, labels: List[str] | Dict[str, Any] | List[List[str]] | List[Dict[str, Any]]
     ) -> List[str] | List[List[str]]:
@@ -405,6 +414,7 @@ def __call__(
         examples: List[Dict[str, Any]] | None = None,
         prompt: str | List[str] | None = None,
         return_hierarchical: bool = False,
+        adapter_ids: str | List[str] | None = None,
     ):
         """
         Perform zero-shot classification.
@@ -421,6 +431,7 @@ def __call__(
             examples: Few-shot examples with 'text' and 'labels'/'true_labels' keys
             prompt: Task description - string (same for all) or list (per-text)
             return_hierarchical: If True, return hierarchical structure with all scores
+            adapter_ids: Optional LoRA adapter id for all texts or one adapter id per text.
 
         Returns:
             List of classification results or hierarchical dict structure.
@@ -430,6 +441,7 @@ def __call__(
         texts = self._normalize_texts(texts)
         thresholds = self._normalize_thresholds(threshold, len(texts))
         classification_types = self._normalize_classification_types(classification_type, len(texts))
+        adapter_ids = self._normalize_adapter_ids(adapter_ids, len(texts))
 
         if rac_examples:
             if len(texts) == 1 and not isinstance(rac_examples[0], list):
@@ -460,12 +472,17 @@ def __call__(
 
             batch_examples = self._get_batch_examples(examples, idx, len(batch_texts))
             batch_prompt = self._get_batch_prompt(prompt, idx, len(batch_texts))
+            batch_adapter_ids = adapter_ids[idx : idx + len(batch_texts)] if adapter_ids is not None else None
 
             tokenized_inputs = self.prepare_inputs(
                 batch_texts, batch_labels, same_labels, examples=batch_examples, prompt=batch_prompt
             )
             max_num_classes = self._resolve_max_num_classes(batch_labels, same_labels)
-            model_output = self.model(**tokenized_inputs, max_num_classes=max_num_classes)
+            model_output = self.model(
+                **tokenized_inputs,
+                max_num_classes=max_num_classes,
+                adapter_ids=batch_adapter_ids,
+            )
             logits = model_output.logits
             probs = torch.sigmoid(logits)
 
@@ -894,6 +911,7 @@ def __call__(
         examples: List[Dict[str, Any]] | None = None,
         prompt: str | List[str] | None = None,
         return_hierarchical: bool = False,
+        adapter_ids: str | List[str] | None = None,
     ):
         """
         Perform zero-shot classification.
@@ -913,6 +931,7 @@ def __call__(
             examples: Few-shot examples, each with 'text' and 'labels' keys
             prompt: Task description - string or list of strings (per-text)
             return_hierarchical: If True, return structure matching input labels
+            adapter_ids: Optional LoRA adapter id for all texts or one adapter id per text.
 
         Returns:
             List of predictions (flat) or hierarchical dicts with all scores
@@ -927,6 +946,7 @@ def __call__(
             examples=examples,
             prompt=prompt,
             return_hierarchical=return_hierarchical,
+            adapter_ids=adapter_ids,
         )