diff --git a/gliclass/__init__.py b/gliclass/__init__.py index 8f9e7e0..ce42569 100644 --- a/gliclass/__init__.py +++ b/gliclass/__init__.py @@ -6,7 +6,7 @@ ZeroShotClassificationWithChunkingPipeline, ) -__version__ = "0.1.18" +__version__ = "0.1.19" # Serve module (optional import) try: diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py index 6dfdfd5..44743f2 100644 --- a/gliclass/pipeline.py +++ b/gliclass/pipeline.py @@ -213,6 +213,42 @@ def __init__( # Ensure model is in evaluation mode for inference self.model.eval() + def _normalize_classification_type(self, classification_type: str | None) -> str: + if classification_type is None: + return self.classification_type + + normalized = classification_type.strip().lower() + if normalized in {"single", "single-label", "single_label"}: + return "single-label" + if normalized in {"multi", "multi-label", "multi_label"}: + return "multi-label" + raise ValueError("Unsupported classification type: choose 'single-label' or 'multi-label'") + + def _normalize_texts(self, texts: str | List[str]) -> List[str]: + if isinstance(texts, str): + return [texts] + return texts + + def _normalize_thresholds(self, threshold: float | List[float], num_texts: int) -> List[float]: + if isinstance(threshold, list): + if len(threshold) != num_texts: + raise ValueError("Length of threshold list must match number of texts.") + return threshold + return [threshold] * num_texts + + def _normalize_classification_types( + self, + classification_type: str | List[str] | None, + num_texts: int, + ) -> List[str]: + if isinstance(classification_type, list): + if len(classification_type) != num_texts: + raise ValueError("Length of classification_type list must match number of texts.") + return [self._normalize_classification_type(item) for item in classification_type] + + normalized = self._normalize_classification_type(classification_type) + return [normalized] * num_texts + def _process_labels( self, labels: List[str] | Dict[str, Any] | List[List[str]] | List[Dict[str, Any]] ) -> List[str] | List[List[str]]: @@ -244,10 +280,29 @@ def _process_labels( def _format_examples_for_input(self, examples: List[Dict[str, Any]] | None = None) -> str: """Format few-shot examples using <> and <> tokens.""" + if not examples: + return "" + examples = [example for example in examples if example is not None] if not examples: return "" return format_examples_prompt(examples, example_token=self.example_token, sep_token=self.sep_token) + def _examples_are_per_text(self, examples) -> bool: + """Detect whether examples are provided per text rather than shared.""" + if not isinstance(examples, list) or len(examples) == 0: + return False + if all(isinstance(example, dict) for example in examples): + return False + return all(example is None or isinstance(example, list) for example in examples) + + def _get_text_examples(self, examples, index: int): + """Get examples for a single text from shared or per-text input.""" + if not examples: + return None + if self._examples_are_per_text(examples): + return examples[index] if index < len(examples) else None + return examples + def _format_prompt(self, prompt: str | List[str] | None = None, index: int = 0) -> str: """Format the task description prompt.""" if prompt is None: @@ -278,7 +333,7 @@ def _get_batch_examples(self, examples, start_idx, batch_size): """Get examples for current batch.""" if not examples: return None - if isinstance(examples[0], list): + if self._examples_are_per_text(examples): return examples[start_idx : start_idx + batch_size] return examples @@ -343,8 +398,9 @@ def __call__( self, texts: str | List[str], labels: List[str] | Dict[str, Any] | List[List[str]] | List[Dict[str, Any]], - threshold: float = 0.5, + threshold: float | List[float] = 0.5, batch_size: int = 8, + classification_type: str | List[str] | None = None, rac_examples: List | None = None, examples: List[Dict[str, Any]] | None = None, prompt: str | List[str] | None = None, @@ -356,8 +412,11 @@ def __call__( Args: texts: Single text or list of texts to classify labels: Labels in various formats (flat list or hierarchical dict) - threshold: Classification threshold for multi-label (default: 0.5) + threshold: Classification threshold for multi-label, either one + value for all texts or one value per text batch_size: Batch size for processing + classification_type: Override classification mode globally or per text. + If None, uses the pipeline's configured classification_type rac_examples: Retrieval augmented examples (legacy) examples: Few-shot examples with 'text' and 'labels'/'true_labels' keys prompt: Task description - string (same for all) or list (per-text) @@ -368,12 +427,15 @@ def __call__( """ original_labels = labels - if isinstance(texts, str): - if rac_examples: - texts = retrieval_augmented_text(texts, rac_examples) - texts = [texts] - elif rac_examples: - texts = [retrieval_augmented_text(text, ex) for text, ex in zip(texts, rac_examples)] + texts = self._normalize_texts(texts) + thresholds = self._normalize_thresholds(threshold, len(texts)) + classification_types = self._normalize_classification_types(classification_type, len(texts)) + + if rac_examples: + if len(texts) == 1 and not isinstance(rac_examples[0], list): + texts = [retrieval_augmented_text(texts[0], rac_examples)] + else: + texts = [retrieval_augmented_text(text, ex) for text, ex in zip(texts, rac_examples)] processed_labels = self._process_labels(labels) @@ -405,14 +467,20 @@ def __call__( max_num_classes = self._resolve_max_num_classes(batch_labels, same_labels) model_output = self.model(**tokenized_inputs, max_num_classes=max_num_classes) logits = model_output.logits + probs = torch.sigmoid(logits) - if self.classification_type == "single-label": - for i in range(len(batch_texts)): - score = torch.softmax(logits[i], dim=-1) - if same_labels: - curr_labels = batch_labels - else: - curr_labels = batch_labels[i] + for i in range(len(batch_texts)): + global_idx = idx + i + item_classification_type = classification_types[global_idx] + item_threshold = thresholds[global_idx] + + if same_labels: + curr_labels = batch_labels + else: + curr_labels = batch_labels[i] + + if item_classification_type == "single-label": + score = torch.softmax(logits[i][: len(curr_labels)], dim=-1) if return_hierarchical: all_scores = {curr_labels[j]: score[j].item() for j in range(len(curr_labels))} @@ -420,16 +488,8 @@ def __call__( pred_label = curr_labels[torch.argmax(score).item()] results.append([{"label": pred_label, "score": score.max().item()}]) - - elif self.classification_type == "multi-label": - sigmoid = torch.nn.Sigmoid() - probs = sigmoid(logits) - for i in range(len(batch_texts)): + elif item_classification_type == "multi-label": text_results = [] - if same_labels: - curr_labels = batch_labels - else: - curr_labels = batch_labels[i] if return_hierarchical: all_scores = {curr_labels[j]: probs[i][j].item() for j in range(len(curr_labels))} @@ -437,11 +497,11 @@ def __call__( for j, prob in enumerate(probs[i][: len(curr_labels)]): score = prob.item() - if score >= threshold: + if score >= item_threshold: text_results.append({"label": curr_labels[j], "score": score}) results.append(text_results) - else: - raise ValueError("Unsupported classification type: choose 'single-label' or 'multi-label'") + else: + raise ValueError("Unsupported classification type: choose 'single-label' or 'multi-label'") if return_hierarchical: hierarchical_results = [] @@ -507,24 +567,12 @@ def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt if same_labels: for i, text in enumerate(texts): - text_examples = None - if examples: - if isinstance(examples[0], list): - text_examples = examples[i] if i < len(examples) else None - else: - text_examples = examples - + text_examples = self._get_text_examples(examples, i) text_prompt = self._format_prompt(prompt, i) inputs.append(self.prepare_input(text, labels, text_examples, text_prompt)) else: for i, (text, labels_) in enumerate(zip(texts, labels)): - text_examples = None - if examples: - if isinstance(examples[0], list): - text_examples = examples[i] if i < len(examples) else None - else: - text_examples = examples - + text_examples = self._get_text_examples(examples, i) text_prompt = self._format_prompt(prompt, i) inputs.append(self.prepare_input(text, labels_, text_examples, text_prompt)) @@ -572,24 +620,14 @@ def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt if same_labels: for i, text in enumerate(texts): - text_examples = None - if examples: - if isinstance(examples[0], list): - text_examples = examples[i] if i < len(examples) else None - else: - text_examples = examples + text_examples = self._get_text_examples(examples, i) text_prompt = self._format_prompt(prompt, i) prompts.append(self.prepare_labels_prompt(labels, text_prompt)) examples_str = self._format_examples_for_input(text_examples) if text_examples else "" processed_texts.append(text + examples_str) else: for i, labels_ in enumerate(labels): - text_examples = None - if examples: - if isinstance(examples[0], list): - text_examples = examples[i] if i < len(examples) else None - else: - text_examples = examples + text_examples = self._get_text_examples(examples, i) text_prompt = self._format_prompt(prompt, i) prompts.append(self.prepare_labels_prompt(labels_, text_prompt)) examples_str = self._format_examples_for_input(text_examples) if text_examples else "" @@ -651,22 +689,12 @@ def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt inputs = [] if same_labels: for i, text in enumerate(texts): - text_examples = None - if examples: - if isinstance(examples[0], list): - text_examples = examples[i] if i < len(examples) else None - else: - text_examples = examples + text_examples = self._get_text_examples(examples, i) text_prompt = self._format_prompt(prompt, i) inputs.append(self.prepare_input(text, labels, text_examples, text_prompt)) else: for i, (text, labels_) in enumerate(zip(texts, labels)): - text_examples = None - if examples: - if isinstance(examples[0], list): - text_examples = examples[i] if i < len(examples) else None - else: - text_examples = examples + text_examples = self._get_text_examples(examples, i) text_prompt = self._format_prompt(prompt, i) inputs.append(self.prepare_input(text, labels_, text_examples, text_prompt)) else: @@ -859,8 +887,9 @@ def __call__( self, texts: str | List[str], labels: List[str] | Dict[str, Any] | List[List[str]] | List[Dict[str, Any]], - threshold: float = 0.5, + threshold: float | List[float] = 0.5, batch_size: int = 8, + classification_type: str | List[str] | None = None, rac_examples: List | None = None, examples: List[Dict[str, Any]] | None = None, prompt: str | List[str] | None = None, @@ -875,8 +904,11 @@ def __call__( Examples: - ["positive", "negative"] - flat labels - {"sentiment": ["positive", "negative"], "topic": ["product", "service"]} - threshold: Classification threshold for multi-label (default: 0.5) + threshold: Classification threshold for multi-label, either one + value for all texts or one value per text batch_size: Batch size for processing + classification_type: Override classification mode globally or per text. + If None, uses the pipeline's configured classification_type rac_examples: Retrieval augmented examples (legacy) examples: Few-shot examples, each with 'text' and 'labels' keys prompt: Task description - string or list of strings (per-text) @@ -890,6 +922,7 @@ def __call__( labels, threshold=threshold, batch_size=batch_size, + classification_type=classification_type, rac_examples=rac_examples, examples=examples, prompt=prompt, @@ -980,22 +1013,12 @@ def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt if same_labels: for i, text in enumerate(texts): - text_examples = None - if examples: - if isinstance(examples[0], list): - text_examples = examples[i] if i < len(examples) else None - else: - text_examples = examples + text_examples = self._get_text_examples(examples, i) text_prompt = self._format_prompt(prompt, i) inputs.append(self.prepare_input(text, labels, text_examples, text_prompt)) else: for i, (text, labels_) in enumerate(zip(texts, labels)): - text_examples = None - if examples: - if isinstance(examples[0], list): - text_examples = examples[i] if i < len(examples) else None - else: - text_examples = examples + text_examples = self._get_text_examples(examples, i) text_prompt = self._format_prompt(prompt, i) inputs.append(self.prepare_input(text, labels_, text_examples, text_prompt)) diff --git a/gliclass/serve/__init__.py b/gliclass/serve/__init__.py index 36bcb35..8728f9c 100644 --- a/gliclass/serve/__init__.py +++ b/gliclass/serve/__init__.py @@ -4,7 +4,6 @@ from .config import GLiClassServeConfig from .memory import GLiClassMemoryEstimator from .server import GLiClassServer, GLiClassFactory, shutdown, serve_gliclass -from .server_model import GLiClassServerModel __all__ = [ "GLiClassClient", @@ -12,7 +11,6 @@ "GLiClassMemoryEstimator", "GLiClassServeConfig", "GLiClassServer", - "GLiClassServerModel", "serve_gliclass", "shutdown", ] diff --git a/gliclass/serve/server.py b/gliclass/serve/server.py index fbd8d7c..28e12fd 100644 --- a/gliclass/serve/server.py +++ b/gliclass/serve/server.py @@ -6,10 +6,13 @@ import torch from ray import serve +from transformers import AutoTokenizer + +from gliclass.model import GLiClassModel +from gliclass.pipeline import ZeroShotClassificationPipeline from .config import GLiClassServeConfig from .memory import GLiClassMemoryEstimator -from .server_model import GLiClassServerModel logger = logging.getLogger("ray.serve") @@ -56,13 +59,23 @@ def __init__(self, config: GLiClassServeConfig): logger.info("Loading model: %s", config.model) - self.model = GLiClassServerModel( - model_name=config.model, - device=self.device, - dtype=self.torch_dtype, - max_length=config.max_model_len, - max_classes=config.max_labels if config.max_labels > 0 else 100, - max_labels_alloc=config.max_labels_alloc, + self.model = GLiClassModel.from_pretrained(config.model) + self.model.config.max_labels_alloc = config.max_labels_alloc + self.model.to(device=self.device, dtype=self.torch_dtype) + self.model.eval() + + self.tokenizer = AutoTokenizer.from_pretrained(config.model) + pipeline_kwargs = { + "model": self.model, + "tokenizer": self.tokenizer, + "max_classes": config.max_labels if config.max_labels > 0 else 100, + "max_length": config.max_model_len, + "device": self.device, + "progress_bar": False, + } + self.pipeline = ZeroShotClassificationPipeline( + classification_type="multi-label", + **pipeline_kwargs, ) if torch.cuda.is_available(): @@ -79,8 +92,8 @@ def __init__(self, config: GLiClassServeConfig): def _precompile(self) -> None: logger.info("Precompiling model for batch sizes: %s", self.config.precompiled_batch_sizes) - if hasattr(self.model.model, "compile"): - self.model.model.compile() + if hasattr(self.model, "compile"): + self.model.compile() dummy_labels = ["person", "organization", "location"] @@ -136,7 +149,7 @@ def batch_size_fn(self, seq_len: int | None = None) -> int: def observed_seq_len( self, texts: list[str], - labels: list[str] | None = None, + labels: list[str] | list[list[str]] | None = None, ) -> int: """Total input word count: longest text + all label words. @@ -153,7 +166,10 @@ def observed_seq_len( max_text_words = max((len(t.split()) for t in texts if t.strip()), default=0) prompt_words = 0 if labels: - prompt_words += sum(len(label.split()) for label in labels) + if isinstance(labels[0], list): + prompt_words += max(sum(len(label.split()) for label in label_set) for label_set in labels) + else: + prompt_words += sum(len(label.split()) for label in labels) total = max_text_words + prompt_words return min(max(total, self.config.calibration_min_seq_len), self.config.max_model_len) @@ -167,34 +183,39 @@ def _filter_labels(self, labels: list[str]) -> list[str]: def _run_batch_internal( self, texts: list[str], - labels: list[str], - threshold: float = 0.5, - multi_label: bool = True, - examples: list[dict[str, Any]] | None = None, + labels: list[str] | list[list[str]], + threshold: float | list[float] = 0.5, + multi_label: bool | list[bool] = True, + examples: list[dict[str, Any]] | list[list[dict[str, Any]] | None] | None = None, prompt: str | list[str] | None = None, - ) -> list[dict[str, Any]]: - """Run batch inference using low-level methods. - - This is the core inference method that uses prepare_batch, tokenize_batch, - collate_batch, run_batch, and decode_batch directly. + ) -> list[list[dict[str, Any]]]: + """Run batch inference using the shared zero-shot pipeline. Args: texts: List of input texts - labels: Label list (same for all texts) - threshold: Classification threshold - multi_label: Multi-label classification - examples: Few-shot examples - prompt: Task description + labels: Shared label list or one label list per text + threshold: Shared threshold or one threshold per text + multi_label: Shared mode or one mode flag per text + examples: Shared examples or one example set per text + prompt: Shared prompt or one prompt per text Returns: List of prediction dicts """ - prepared = self.model.prepare_batch(texts, labels, examples, prompt) - tokenized = self.model.tokenize_batch(prepared) - batch = self.model.collate_batch(tokenized) - outputs = self.model.run_batch(batch) - results = self.model.decode_batch(outputs, threshold, multi_label) - return results + if isinstance(multi_label, list): + classification_type = ["multi-label" if item else "single-label" for item in multi_label] + else: + classification_type = "multi-label" if multi_label else "single-label" + + return self.pipeline( + texts, + labels, + threshold=threshold, + batch_size=max(len(texts), 1), + classification_type=classification_type, + examples=examples, + prompt=prompt, + ) def predict( self, @@ -204,7 +225,7 @@ def predict( multi_label: bool = True, examples: list[dict[str, Any]] | None = None, prompt: str | list[str] | None = None, - ) -> list[dict[str, Any]]: + ) -> list[list[dict[str, Any]]]: if isinstance(texts, str): texts = [texts] @@ -267,33 +288,27 @@ async def _infer_batch( using batch_size_fn on the observed seq length — so the next accumulation picks the largest precompiled size that fits. - Assumes batch requests are homogeneous — labels/thresholds/flags - are taken from the first request. + Supports heterogeneous request parameters by passing per-text + thresholds, classification types, labels, examples, and prompts + through to the shared pipeline. """ # Dynamically adjust batch size based on observed sequence length next_max_batch = self.server.batch_size_fn( seq_len=self.server.observed_seq_len( texts, - labels=labels_list[0] if labels_list else None, + labels=labels_list, ) ) self._infer_batch.set_max_batch_size(next_max_batch) - # Use first request's parameters (assume homogeneous) - labels = labels_list[0] - threshold = thresholds[0] - multi_label = multi_label_list[0] - examples = examples_list[0] - prompt = prompts_list[0] - # Process entire batch at once results = self.server._run_batch_internal( texts, - labels, - threshold=threshold, - multi_label=multi_label, - examples=examples, - prompt=prompt, + labels_list, + threshold=thresholds, + multi_label=multi_label_list, + examples=examples_list, + prompt=prompts_list, ) return results diff --git a/gliclass/serve/server_model.py b/gliclass/serve/server_model.py deleted file mode 100644 index b93955c..0000000 --- a/gliclass/serve/server_model.py +++ /dev/null @@ -1,336 +0,0 @@ -"""Server-optimized model wrapper for GLiClass with low-level batch processing.""" - -import logging -from typing import Any - -import torch -from transformers import AutoTokenizer - -from gliclass.model import GLiClassModel -from gliclass.pipeline import format_examples_prompt - -logger = logging.getLogger(__name__) - - -class GLiClassServerModel: - """Server-optimized wrapper combining pipeline logic and model. - - Splits inference into low-level stages for CPU/GPU parallelism: - - prepare_batch (CPU): text formatting, label flattening - - tokenize_batch (CPU): tokenization, input creation - - collate_batch (CPU): tensor creation, padding - - run_batch (GPU): forward pass, encoding - - decode_batch (CPU/GPU): sigmoid, threshold filtering, result formatting - - Avoids DataLoader overhead by keeping tokenizers/collators alive - and calling model methods directly. - """ - - def __init__( - self, - model_name: str, - device: torch.device, - dtype: torch.dtype, - max_length: int = 2048, - max_classes: int = 25, - max_labels_alloc: str | int = "dynamic", # "dynamic" | int | "fixed" - ): - self.device = device - self.dtype = dtype - self.max_length = max_length - self.max_classes = max_classes - self.max_labels_alloc = max_labels_alloc - - logger.info(f"Loading model: {model_name}") - self.model = GLiClassModel.from_pretrained(model_name) - self.model.config.max_labels_alloc = max_labels_alloc - self.model.to(device=device, dtype=dtype) - self.model.eval() - - self.architecture_type = self.model.config.architecture_type - self.prompt_first = getattr(self.model.config, "prompt_first", False) - - logger.info(f"Loading tokenizers for {self.architecture_type}") - # Use model_name, not encoder_model_name, to get tokenizer with special tokens - self.tokenizer = AutoTokenizer.from_pretrained(model_name) - - if self.architecture_type in {"bi-encoder", "bi-encoder-fused"}: - # For bi-encoder, load labels tokenizer from model (not label_model_name) - # to ensure special tokens are available - self.labels_tokenizer = AutoTokenizer.from_pretrained(model_name) - else: - self.labels_tokenizer = None - - self.label_token = "<