diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..a271395 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,70 @@ +name: Tests + +on: + push: + branches: + - main + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + name: pytest (Python ${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12"] + + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Cache pip + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-py${{ matrix.python-version }}-pip-${{ hashFiles('pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-py${{ matrix.python-version }}-pip- + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e . + pip install pytest pytest-asyncio + + - name: Run pytest + run: pytest -v --tb=short + + lint: + name: ruff + runs-on: ubuntu-latest + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install ruff + run: pip install ruff + + - name: ruff check + run: ruff check gliclass + + - name: ruff format --check + run: ruff format --check gliclass diff --git a/.gitignore b/.gitignore index c59cda0..0a9ee60 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ gradio_cached_examples/ test.ipynb demo1.py .gradio/ +uv.lock # Distribution / packaging .Python @@ -57,6 +58,7 @@ coverage.xml *.py,cover .hypothesis/ .pytest_cache/ +.ruff_cache/ cover/ # Translations diff --git a/README.md b/README.md index 00d4f88..1a829cb 100644 --- a/README.md +++ b/README.md @@ -196,6 +196,64 @@ for predict in results: print(f"{predict['label']} => {predict['score']:.3f}") ``` +### 🚀 Production Serving + +Deploy GLiClass with Ray Serve for production workloads with dynamic batching and memory-aware processing. + +#### Installation + +```bash +pip install gliclass[serve] +``` + +#### Quick Start + +```bash +# Default model +python -m gliclass.serve + +# Specify model and port +python -m gliclass.serve --model knowledgator/gliclass-edge-v3.0 --port 8000 + +# With config file +python -m gliclass.serve --config serve_configs/serve_config.yaml +``` + +#### Python Client + +```python +from gliclass.serve import GLiClassClient + +client = GLiClassClient(url="http://localhost:8000/gliclass") + +result = client.classify( + text="This is a great product!", + labels=["positive", "negative", "neutral"], + threshold=0.3, +) +print(result) # [{"label": "positive", "score": 0.95}, ...] +``` + +#### HTTP API + +The HTTP endpoint processes one text per request. + +```bash +curl -X POST http://localhost:8000/gliclass \ + -H "Content-Type: application/json" \ + -d '{ + "texts": "This is a great product!", + "labels": ["positive", "negative", "neutral"], + "threshold": 0.3 + }' + +# Response: [{"label": "positive", "score": 0.95}, ...] +``` + +**Note:** For batch processing multiple texts, use the `ZeroShotClassificationPipeline` directly instead of the serving API. + +See `serve_configs/serve_config.yaml` for full configuration options. + ### 🎯 Key Use Cases - **Sentiment Analysis:** Rapidly classify texts as positive, negative, or neutral. @@ -281,8 +339,6 @@ config = GLiClassModelConfig( ) ``` -Gotcha — here’s a **much leaner, README-style version**, no fluff, just what matters 👇 - --- ### Flash Attention Backends diff --git a/gliclass/__init__.py b/gliclass/__init__.py index a548992..8f9e7e0 100644 --- a/gliclass/__init__.py +++ b/gliclass/__init__.py @@ -1,5 +1,15 @@ from .model import GLiClassModel, GLiClassBiEncoder, GLiClassUniEncoder, GLiClassEncoderDecoderCLS from .config import GLiClassModelConfig -from .pipeline import ZeroShotClassificationPipeline, BiEncoderZeroShotClassificationPipeline, ZeroShotClassificationWithChunkingPipeline +from .pipeline import ( + ZeroShotClassificationPipeline, + BiEncoderZeroShotClassificationPipeline, + ZeroShotClassificationWithChunkingPipeline, +) -__version__ = "0.1.18" \ No newline at end of file +__version__ = "0.1.18" + +# Serve module (optional import) +try: + from . import serve +except ImportError: + serve = None diff --git a/gliclass/config.py b/gliclass/config.py index 69f37ba..8698977 100644 --- a/gliclass/config.py +++ b/gliclass/config.py @@ -1,11 +1,11 @@ from transformers import AutoConfig -from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging from transformers.models.auto import CONFIG_MAPPING +from transformers.configuration_utils import PretrainedConfig from .utils import is_module_available -IS_TURBOT5 = is_module_available('turbot5') +IS_TURBOT5 = is_module_available("turbot5") if IS_TURBOT5: from turbot5.model.config import T5Config @@ -22,26 +22,26 @@ class GLiClassModelConfig(PretrainedConfig): def __init__( self, - encoder_config = None, + encoder_config=None, encoder_model=None, label_model_config=None, label_model_name=None, - class_token_index = -1, - text_token_index = -1, - example_token_index = -1, + class_token_index=-1, + text_token_index=-1, + example_token_index=-1, ignore_index=-100, hidden_size=None, projector_hidden_act="gelu", vocab_size=None, - problem_type='single_label_classification', + problem_type="single_label_classification", max_num_classes=25, use_lstm=False, initializer_range=0.03, - scorer_type='simple', + scorer_type="simple", scorer_num_heads=16, scorer_mlp_hidden_size=1024, scorer_attn_dropout=0.1, - pooling_strategy='first', + pooling_strategy="first", class_token_pooling="first", focal_loss_alpha=0.5, focal_loss_gamma=2, @@ -49,23 +49,21 @@ def __init__( logit_scale_init_value=2.6592, normalize_features=False, extract_text_features=False, - max_labels_alloc: str = 'dynamic', + max_labels_alloc: str = "dynamic", contrastive_loss_coef=0, - architecture_type = 'uni-encoder', - prompt_first = False, - squeeze_layers = False, - layer_wise = False, - encoder_layer_id = -1, - embed_class_token = True, - dropout = 0.1, - use_segment_embeddings = False, + architecture_type="uni-encoder", + prompt_first=False, + squeeze_layers=False, + layer_wise=False, + encoder_layer_id=-1, + embed_class_token=True, + dropout=0.1, + use_segment_embeddings=False, **kwargs, ): if isinstance(encoder_config, dict): - encoder_config["model_type"] = (encoder_config["model_type"] - if "model_type" in encoder_config - else "deberta-v2") - if encoder_config['model_type'] == 't5': + encoder_config["model_type"] = encoder_config.get("model_type", "deberta-v2") + if encoder_config["model_type"] == "t5": encoder_config = T5Config(**encoder_config) elif encoder_config["model_type"] in CONFIG_MAPPING: encoder_config = CONFIG_MAPPING[encoder_config["model_type"]](**encoder_config) @@ -83,9 +81,7 @@ def __init__( if label_model_name is not None: if isinstance(label_model_config, dict): - label_model_config["model_type"] = (label_model_config["model_type"] - if "model_type" in label_model_config - else "deberta-v2") + label_model_config["model_type"] = label_model_config.get("model_type", "deberta-v2") label_model_config = CONFIG_MAPPING[label_model_config["model_type"]](**label_model_config) elif label_model_config is None: label_model_config = CONFIG_MAPPING["deberta-v2"]() @@ -104,19 +100,19 @@ def __init__( self.vocab_size = self.encoder_config.vocab_size else: self.vocab_size = vocab_size - + if class_token_index == -1: self.class_token_index = self.vocab_size else: self.class_token_index = class_token_index - + if text_token_index == -1: - self.text_token_index = self.vocab_size+1 + self.text_token_index = self.vocab_size + 1 else: self.text_token_index = text_token_index if example_token_index == -1: - self.example_token_index = self.vocab_size+2 + self.example_token_index = self.vocab_size + 2 else: self.example_token_index = example_token_index @@ -124,20 +120,20 @@ def __init__( self.projector_hidden_act = projector_hidden_act self.problem_type = problem_type self.max_num_classes = max_num_classes - self.initializer_range=initializer_range + self.initializer_range = initializer_range self.scorer_type = scorer_type self.scorer_num_heads = scorer_num_heads self.scorer_mlp_hidden_size = scorer_mlp_hidden_size self.scorer_attn_dropout = scorer_attn_dropout - self.pooling_strategy=pooling_strategy - self.class_token_pooling=class_token_pooling + self.pooling_strategy = pooling_strategy + self.class_token_pooling = class_token_pooling self.use_lstm = use_lstm - self.focal_loss_alpha=focal_loss_alpha - self.focal_loss_gamma=focal_loss_gamma + self.focal_loss_alpha = focal_loss_alpha + self.focal_loss_gamma = focal_loss_gamma self.focal_loss_reduction = focal_loss_reduction - self.contrastive_loss_coef=contrastive_loss_coef + self.contrastive_loss_coef = contrastive_loss_coef self.logit_scale_init_value = logit_scale_init_value - self.normalize_features=normalize_features + self.normalize_features = normalize_features self.extract_text_features = extract_text_features self.max_labels_alloc = max_labels_alloc self.architecture_type = architecture_type @@ -150,4 +146,3 @@ def __init__( self.dropout = dropout self.use_segment_embeddings = use_segment_embeddings super().__init__(**kwargs) - diff --git a/gliclass/data_processing.py b/gliclass/data_processing.py index 370ef2b..8cc361e 100644 --- a/gliclass/data_processing.py +++ b/gliclass/data_processing.py @@ -1,15 +1,18 @@ -import random -import torch import copy +import random from dataclasses import dataclass -from torch.nn.utils.rnn import pad_sequence + +import torch from torch.utils.data import Dataset +from torch.nn.utils.rnn import pad_sequence + @dataclass class AugmentationConfig: """Configuration for data augmentation.""" + enabled: bool = True - + # Probability for each augmentation type random_label_removal_prob: float = 0.15 random_label_addition_prob: float = 0.10 @@ -19,6 +22,7 @@ class AugmentationConfig: random_add_examples_prob: float = 0.25 max_num_examples: int = 5 + class DataAugmenter: def __init__(self, config, examples, labels, label2description=None): self.config = config @@ -26,9 +30,9 @@ def __init__(self, config, examples, labels, label2description=None): self.labels = sorted(labels) self.max_examples = self.config.max_num_examples self.label2description = label2description or {} - + def remove_labels(self, true_labels, all_labels): - if len(all_labels)<=1: + if len(all_labels) <= 1: return true_labels, all_labels k = random.randint(1, len(all_labels)) all_labels = random.sample(all_labels, k=k) @@ -43,159 +47,157 @@ def add_random_labels(self, all_labels): add_labels = random.sample(self.labels, k=k) all_labels.extend(add_labels) return all_labels - + def add_random_text(self, text, all_labels): if not self.examples: return text example = random.sample(self.examples, k=1)[0] - curr_labels = example['all_labels'] + curr_labels = example["all_labels"] joint_labels = set(all_labels) & set(curr_labels) if len(joint_labels): return text else: if random.randint(0, 1): - text = example['text'] + ' ' + text + text = example["text"] + " " + text else: - text = text + ' ' + example['text'] + text = text + " " + example["text"] return text - + def add_random_synonyms(self, all_labels): """Replace some labels with their synonyms if available.""" if not self.label2description: return all_labels - + augmented_labels = [] for label in all_labels: if label in self.label2description: label_info = self.label2description[label] - synonyms = label_info.get('synonyms', []) - + synonyms = label_info.get("synonyms", []) + if synonyms and random.random() < 0.5: augmented_labels.append(random.choice(synonyms)) else: augmented_labels.append(label) else: augmented_labels.append(label) - + return augmented_labels - + def add_random_descriptions(self, item): """Add descriptions to labels in the text or metadata.""" - if not self.label2description or not item['all_labels']: + if not self.label2description or not item["all_labels"]: return item - - max_labels = min(3, len(item['all_labels'])) - labels_to_describe = random.sample( - item['all_labels'], - k=random.randint(1, max_labels) - ) - + + max_labels = min(3, len(item["all_labels"])) + labels_to_describe = random.sample(item["all_labels"], k=random.randint(1, max_labels)) + descriptions = [] for label in labels_to_describe: if label in self.label2description: label_info = self.label2description[label] - desc_list = label_info.get('descriptions', []) + desc_list = label_info.get("descriptions", []) if desc_list: descriptions.append(f"{label}: {random.choice(desc_list)}") - + if descriptions: - desc_text = ' '.join(descriptions) + desc_text = " ".join(descriptions) if random.random() < 0.5: - item['text'] = desc_text + ' ' + item['text'] + item["text"] = desc_text + " " + item["text"] else: - item['text'] = item['text'] + ' ' + desc_text - + item["text"] = item["text"] + " " + desc_text + return item - + def add_random_examples(self, item): """Add example texts with similar labels.""" - if not item['all_labels']: + if not item["all_labels"]: return item - + candidate_examples = item.get("examples", []) - item_label_set = set(item['all_labels']) + item_label_set = set(item["all_labels"]) if not candidate_examples: - for example in self.examples: - example_label_set = set(example['true_labels']) - example_text = example['text'] + example_label_set = set(example["true_labels"]) + example_text = example["text"] overlap = item_label_set & example_label_set - + # Only consider examples with at least one overlapping label if overlap: candidate_examples.append({"text": example_text, "labels": list(example_label_set)}) - + if not candidate_examples: return item - + # Sort by overlap and select top examples random.shuffle(candidate_examples) - top_candidates = candidate_examples[:self.max_examples] + top_candidates = candidate_examples[: self.max_examples] num_examples = random.randint(1, min(2, len(top_candidates))) selected_examples = random.sample(top_candidates, k=num_examples) - - item['examples'] = selected_examples - + + item["examples"] = selected_examples + return item - + def augment(self, item): if not self.config.enabled: return item - text = copy.deepcopy(item['text']) - true_labels = copy.deepcopy(item['true_labels']) - all_labels = copy.deepcopy(item['all_labels']) + text = copy.deepcopy(item["text"]) + true_labels = copy.deepcopy(item["true_labels"]) + all_labels = copy.deepcopy(item["all_labels"]) # Create augmented item - aug_item = { - 'text': text, - 'true_labels': true_labels, - 'all_labels': all_labels - } - + aug_item = {"text": text, "true_labels": true_labels, "all_labels": all_labels} + # Copy any additional fields for key in item: if key not in aug_item: aug_item[key] = copy.deepcopy(item[key]) if random.random() < self.config.random_label_removal_prob: - aug_item['true_labels'], aug_item['all_labels'] = self.remove_labels( - aug_item['true_labels'], aug_item['all_labels'] + aug_item["true_labels"], aug_item["all_labels"] = self.remove_labels( + aug_item["true_labels"], aug_item["all_labels"] ) - + if random.random() < self.config.random_label_addition_prob: - aug_item['all_labels'] = self.add_random_labels(aug_item['all_labels']) + aug_item["all_labels"] = self.add_random_labels(aug_item["all_labels"]) if random.random() < self.config.random_text_addition_prob: - aug_item['text'] = self.add_random_text(aug_item['text'], aug_item['all_labels']) - + aug_item["text"] = self.add_random_text(aug_item["text"], aug_item["all_labels"]) + if random.random() < self.config.random_add_synonyms_prob: - aug_item['all_labels'] = self.add_random_synonyms(aug_item['all_labels']) - + aug_item["all_labels"] = self.add_random_synonyms(aug_item["all_labels"]) + if random.random() < self.config.random_add_description_prob: aug_item = self.add_random_descriptions(aug_item) - + if random.random() < self.config.random_add_examples_prob: aug_item = self.add_random_examples(aug_item) return aug_item - + + class GLiClassDataset(Dataset): - def __init__(self, examples, tokenizer, augment_config, - label2description={}, - max_length=512, - problem_type='multi_label_classification', - architecture_type = 'uni-encoder', - add_description=True, - prompt_first=False, - get_negatives = False, - max_labels = 50, - labels_tokenizer=None, - shuffle_labels = True): + def __init__( + self, + examples, + tokenizer, + augment_config, + label2description={}, + max_length=512, + problem_type="multi_label_classification", + architecture_type="uni-encoder", + add_description=True, + prompt_first=False, + get_negatives=False, + max_labels=50, + labels_tokenizer=None, + shuffle_labels=True, + ): self.tokenizer = tokenizer self.labels_tokenizer = labels_tokenizer self.label2description = label2description @@ -210,44 +212,44 @@ def __init__(self, examples, tokenizer, augment_config, self.get_negatives = get_negatives self.max_labels = max_labels self.shuffle_labels = shuffle_labels - + self.sep_token = "<>" self.label_token = "<