[ { "path": ".github/workflows/release.yaml", "content": "name: Release GLiClass to PyPI\n\non:\n push:\n tags:\n - 'v*' # Trigger on version tags (e.g., v1.0.0, v2.1.3)\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.ref }}\n cancel-in-progress: true\n\njobs:\n build:\n name: Build distribution πŸ“¦\n runs-on: ubuntu-latest\n\n steps:\n - uses: actions/checkout@v6\n with:\n persist-credentials: false\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: \"3.x\"\n - name: Install pypa/build\n run: >-\n python3 -m\n pip install\n build\n --user\n - name: Build a binary wheel and a source tarball\n run: python3 -m build\n - name: Store the distribution packages\n uses: actions/upload-artifact@v5\n with:\n name: python-package-distributions\n path: dist/\n\n publish-to-pypi:\n name: >-\n Publish Python 🐍 distribution πŸ“¦ to PyPI\n if: startsWith(github.ref, 'refs/tags/') # only publish to PyPI on tag pushes\n needs:\n - build\n runs-on: ubuntu-latest\n environment:\n name: pypi\n url: https://pypi.org/project/gliclass/ # Replace with your PyPI project name\n permissions:\n id-token: write # IMPORTANT: mandatory for trusted publishing\n\n steps:\n - name: Checkout code\n uses: actions/checkout@v6\n with:\n fetch-depth: 0 # Fetch all history to check branches\n - name: Verify tag is on main branch\n run: |\n if ! git branch -r --contains ${{ github.ref_name }} | grep -q 'origin/main'; then\n echo \"Error: Tag ${{ github.ref_name }} is not on the main branch\"\n exit 1\n fi\n echo \"βœ“ Tag ${{ github.ref_name }} is on main branch\"\n - name: Download all the dists\n uses: actions/download-artifact@v6\n with:\n name: python-package-distributions\n path: dist/\n - name: Publish distribution πŸ“¦ to PyPI\n uses: pypa/gh-action-pypi-publish@release/v1" }, { "path": ".github/workflows/tests.yml", "content": "name: Tests\n\non:\n push:\n branches:\n - main\n pull_request:\n branches:\n - main\n workflow_dispatch:\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.ref }}\n cancel-in-progress: true\n\njobs:\n test:\n name: pytest (Python ${{ matrix.python-version }})\n runs-on: ubuntu-latest\n strategy:\n fail-fast: false\n matrix:\n python-version: [\"3.10\", \"3.11\", \"3.12\"]\n\n steps:\n - name: Check out repository\n uses: actions/checkout@v4\n\n - name: Set up Python\n uses: actions/setup-python@v5\n with:\n python-version: ${{ matrix.python-version }}\n\n - name: Cache pip\n uses: actions/cache@v4\n with:\n path: ~/.cache/pip\n key: ${{ runner.os }}-py${{ matrix.python-version }}-pip-${{ hashFiles('pyproject.toml') }}\n restore-keys: |\n ${{ runner.os }}-py${{ matrix.python-version }}-pip-\n\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install -e .\n pip install pytest pytest-asyncio\n\n - name: Run pytest\n run: pytest -v --tb=short\n\n lint:\n name: ruff\n runs-on: ubuntu-latest\n steps:\n - name: Check out repository\n uses: actions/checkout@v4\n\n - name: Set up Python\n uses: actions/setup-python@v5\n with:\n python-version: \"3.12\"\n\n - name: Install ruff\n run: pip install ruff\n\n - name: ruff check\n run: ruff check gliclass\n\n - name: ruff format --check\n run: ruff format --check gliclass\n" }, { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n#custom\nmodels/\nwandb/\ngradio_cached_examples/\ntest.ipynb\ndemo1.py\n.gradio/\nuv.lock\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n.ruff_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# poetry\n# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control\n#poetry.lock\n\n# pdm\n# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.\n#pdm.lock\n# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it\n# in version control.\n# https://pdm.fming.dev/latest/usage/project/#working-with-version-control\n.pdm.toml\n.pdm-python\n.pdm-build/\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n\n# PyCharm\n# JetBrains specific template is maintained in a separate JetBrains.gitignore that can\n# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore\n# and can be added to the global gitignore or merged into this file. For a more nuclear\n# option (not recommended) you can uncomment the following to ignore the entire idea folder.\n#.idea/" }, { "path": "LICENSE", "content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "README.md", "content": "# ⭐ GLiClass: Generalist and Lightweight Model for Sequence Classification\n\n**GLiClass** is an efficient, zero-shot sequence classification model inspired by the [GLiNER](https://github.com/urchade/GLiNER/tree/main) framework. It achieves comparable performance to traditional cross-encoder models while being significantly more computationally efficient, offering classification results approximately **10 times faster** by performing classification in a single forward pass.\n\n

\n πŸ“„ Blog\n   β€’  \n πŸ“’ Discord\n   β€’  \n πŸ“Ί Demo\n   β€’  \n πŸ€— Available models\n   β€’  \n \n \n \n

\n\n### πŸš€ Quick Start\n\nInstall GLiClass easily using pip:\n\n```bash\npip install gliclass\n```\n\n#### Install from Source\n\nClone and install directly from GitHub:\n\n```bash\ngit clone https://github.com/Knowledgator/GLiClass\ncd GLiClass\n\npython -m venv venv\nsource venv/bin/activate # Windows: venv\\Scripts\\activate\n\npip install -r requirements.txt\npip install .\n```\n\nVerify your installation:\n\n```python\nimport gliclass\nprint(gliclass.__version__)\n```\n\n### πŸ§‘β€πŸ’» Usage Example\n\n```python\nfrom gliclass import GLiClassModel, ZeroShotClassificationPipeline\nfrom transformers import AutoTokenizer\n\nmodel = GLiClassModel.from_pretrained(\"knowledgator/gliclass-small-v1.0\")\ntokenizer = AutoTokenizer.from_pretrained(\"knowledgator/gliclass-small-v1.0\")\n\npipeline = ZeroShotClassificationPipeline(\n model, tokenizer, classification_type='multi-label', device='cuda:0'\n)\n\ntext = \"One day I will see the world!\"\nlabels = [\"travel\", \"dreams\", \"sport\", \"science\", \"politics\"]\nresults = pipeline(text, labels, threshold=0.5)[0]\n\nfor result in results:\n print(f\"{result['label']} => {result['score']:.3f}\")\n```\n\n### πŸ”₯ New Features\n\n#### Hierarchical Labels\n\nGLiClass now supports hierarchical label structures using dot notation:\n\n```python\nhierarchical_labels = {\n \"sentiment\": [\"positive\", \"negative\", \"neutral\"],\n \"topic\": [\"product\", \"service\", \"shipping\"]\n}\n\ntext = \"The product quality is amazing but delivery was slow\"\nresults = pipeline(text, hierarchical_labels, threshold=0.5)[0]\n\nfor result in results:\n print(f\"{result['label']} => {result['score']:.3f}\")\n# Output:\n# sentiment.positive => 0.892\n# topic.product => 0.921\n# topic.shipping => 0.763\n```\n\nGet hierarchical output matching your input structure:\n\n```python\nresults = pipeline(text, hierarchical_labels, return_hierarchical=True)[0]\nprint(results)\n# Output:\n# {\n# \"sentiment\": {\"positive\": 0.892, \"negative\": 0.051, \"neutral\": 0.124},\n# \"topic\": {\"product\": 0.921, \"service\": 0.153, \"shipping\": 0.763}\n# }\n```\n\n#### Few-Shot Examples\n\nImprove classification accuracy with in-context examples using the `<>` token:\n\n```python\nexamples = [\n {\n \"text\": \"Love this item, great quality!\",\n \"labels\": [\"positive\", \"product\"]\n },\n {\n \"text\": \"Customer support was unhelpful\",\n \"labels\": [\"negative\", \"service\"]\n }\n]\n\ntext = \"Fast delivery and the item works perfectly!\"\nlabels = [\"positive\", \"negative\", \"product\", \"service\", \"shipping\"]\n\nresults = pipeline(text, labels, examples=examples, threshold=0.5)[0]\n\nfor result in results:\n print(f\"{result['label']} => {result['score']:.3f}\")\n```\n\n#### Task Description Prompts\n\nAdd custom prompts to guide the classification task:\n\n```python\ntext = \"The battery life on this phone is incredible\"\nlabels = [\"positive\", \"negative\", \"neutral\"]\n\nresults = pipeline(\n text,\n labels,\n prompt=\"Classify the sentiment of this product review:\",\n threshold=0.5\n)[0]\n```\n\nUse per-text prompts for batch processing:\n\n```python\ntexts = [\"Review about electronics\", \"Review about clothing\"]\nprompts = [\n \"Analyze this electronics review:\",\n \"Analyze this clothing review:\"\n]\n\nresults = pipeline(texts, labels, prompt=prompts)\n```\n\n#### Long Document Classification\n\nProcess long documents with automatic text chunking:\n\n```python\nfrom gliclass import ZeroShotClassificationWithChunkingPipeline\n\nchunking_pipeline = ZeroShotClassificationWithChunkingPipeline(\n model,\n tokenizer,\n text_chunk_size=8192,\n text_chunk_overlap=256,\n labels_chunk_size=8\n)\n\nlong_document = \"...\" # Very long text\nlabels = [\"category1\", \"category2\", \"category3\"]\n\nresults = chunking_pipeline(long_document, labels, threshold=0.5)\n```\n\n### 🌟 Retrieval-Augmented Classification (RAC)\n\nWith new models trained with retrieval-agumented classification, such as [this model](https://huggingface.co/knowledgator/gliclass-base-v2.0-rac-init) you can specify examples to improve classification accuracy:\n\n```python\nexample = {\n \"text\": \"A new machine learning platform automates complex data workflows but faces integration issues.\",\n \"all_labels\": [\"AI\", \"automation\", \"data_analysis\", \"usability\", \"integration\"],\n \"true_labels\": [\"AI\", \"integration\", \"automation\"]\n}\n\ntext = \"The new AI-powered tool streamlines data analysis but has limited integration capabilities.\"\nlabels = [\"AI\", \"automation\", \"data_analysis\", \"usability\", \"integration\"]\n\nresults = pipeline(text, labels, threshold=0.1, rac_examples=[example])[0]\n\nfor predict in results:\n print(f\"{predict['label']} => {predict['score']:.3f}\")\n```\n\n### πŸš€ Production Serving\n\nDeploy GLiClass with Ray Serve for production workloads with dynamic batching and memory-aware processing.\n\n#### Installation\n\n```bash\npip install gliclass[serve]\n```\n\n#### Quick Start\n\n```bash\n# Default model\npython -m gliclass.serve\n\n# Specify model and port\npython -m gliclass.serve --model knowledgator/gliclass-edge-v3.0 --port 8000\n\n# With config file\npython -m gliclass.serve --config serve_configs/serve_config.yaml\n```\n\n#### Python Client\n\n```python\nfrom gliclass.serve import GLiClassClient\n\nclient = GLiClassClient(url=\"http://localhost:8000/gliclass\")\n\nresult = client.classify(\n text=\"This is a great product!\",\n labels=[\"positive\", \"negative\", \"neutral\"],\n threshold=0.3,\n)\nprint(result) # [{\"label\": \"positive\", \"score\": 0.95}, ...]\n```\n\n#### HTTP API\n\nThe HTTP endpoint processes one text per request.\n\n```bash\ncurl -X POST http://localhost:8000/gliclass \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"texts\": \"This is a great product!\",\n \"labels\": [\"positive\", \"negative\", \"neutral\"],\n \"threshold\": 0.3\n }'\n\n# Response: [{\"label\": \"positive\", \"score\": 0.95}, ...]\n```\n\n**Note:** For batch processing multiple texts, use the `ZeroShotClassificationPipeline` directly instead of the serving API.\n\nSee `serve_configs/serve_config.yaml` for full configuration options.\n\n### 🎯 Key Use Cases\n\n- **Sentiment Analysis:** Rapidly classify texts as positive, negative, or neutral.\n- **Document Classification:** Efficiently organize and categorize large document collections.\n- **Search Results Re-ranking:** Improve relevance and precision by reranking search outputs.\n- **News Categorization:** Automatically tag and organize news articles into predefined categories.\n- **Fact Checking:** Quickly validate and categorize statements based on factual accuracy.\n\n### πŸ› οΈ How to Train\n\nPrepare your training data as follows:\n\n```json\n[\n {\"text\": \"Sample text.\", \"all_labels\": [\"sports\", \"science\", \"business\"], \"true_labels\": [\"sports\"]},\n ...\n]\n```\n\nOptionally, specify confidence scores explicitly:\n\n```json\n[\n {\"text\": \"Sample text.\", \"all_labels\": [\"sports\", \"science\"], \"true_labels\": {\"sports\": 0.9}},\n ...\n]\n```\n\nPlease, refer to the `train.py` script to set up your training from scratch or fine-tune existing models.\n\n### βš™οΈ Advanced Configuration\n\n#### Architecture Types\n\nGLiClass supports multiple architecture types:\n\n- **uni-encoder**: Single encoder for both text and labels (default, most efficient)\n- **bi-encoder**: Separate encoders for text and labels\n- **bi-encoder-fused**: Bi-encoder with label embeddings fused into text encoding\n- **encoder-decoder**: Encoder-decoder architecture for sequence-to-sequence tasks\n\n```python\nfrom gliclass import GLiClassBiEncoder\n\n# Load a bi-encoder model\nmodel = GLiClassBiEncoder.from_pretrained(\"knowledgator/gliclass-biencoder-v1.0\")\n```\n\n#### Pooling Strategies\n\nConfigure how token embeddings are pooled:\n\n- `first`: First token (CLS token)\n- `avg`: Average pooling\n- `max`: Max pooling\n- `last`: Last token\n- `sum`: Sum pooling\n- `rms`: Root mean square pooling\n- `abs_max`: Max of absolute values\n- `abs_avg`: Average of absolute values\n\n```python\nfrom gliclass import GLiClassModelConfig\n\nconfig = GLiClassModelConfig(\n pooling_strategy='avg',\n class_token_pooling='average' # or 'first'\n)\n```\n\n#### Scoring Mechanisms\n\nChoose different scoring mechanisms for classification:\n\n- `simple`: Dot product (fastest)\n- `weighted-dot`: Weighted dot product with learned projections\n- `mlp`: Multi-layer perceptron scorer\n- `hopfield`: Hopfield network-based scorer\n\n```python\nconfig = GLiClassModelConfig(\n scorer_type='mlp'\n)\n```\n\n---\n\n### Flash Attention Backends\n\nGLiClass supports optional flash attention backends for faster inference.\n\n#### Install\n\n```bash\npip install flashdeberta # DeBERTa v2\npip install turbot5 # T5 / mT5\n```\n\n---\n\n#### FlashDeBERTa (DeBERTa v2)\n\nEnable via environment variable:\n\n```bash\nexport USE_FLASHDEBERTA=1\n```\n\nIf `flashdeberta` is installed, DeBERTa v2 models will use `FlashDebertaV2Model`.\nOtherwise, GLiClass falls back to `DebertaV2Model`.\n\n---\n\n#### TurboT5 (T5 / mT5)\n\nEnable via environment variable:\n\n```bash\nexport TURBOT5_ATTN_TYPE=triton-basic\n```\n\nIf `turbot5` is installed, T5 / mT5 models will use `FlashT5EncoderModel`.\nOtherwise, GLiClass falls back to `T5EncoderModel`.\n\nNotes:\n* Flash backends are **optional**\n* Enabled automatically when available\n* No code changes required\n\nWant it even tighter (single block), or is this the sweet spot?\n\n\n## πŸ“š Citations\n\nIf you find GLiClass useful in your research or project, please cite our papers:\n\n\n```bibtex\n@misc{stepanov2025gliclassgeneralistlightweightmodel,\n title={GLiClass: Generalist Lightweight Model for Sequence Classification Tasks}, \n author={Ihor Stepanov and Mykhailo Shtopko and Dmytro Vodianytskyi and Oleksandr Lukashov and Alexander Yavorskyi and Mykyta Yaroshenko},\n year={2025},\n eprint={2508.07662},\n archivePrefix={arXiv},\n primaryClass={cs.LG},\n url={https://arxiv.org/abs/2508.07662}, \n}\n```\n" }, { "path": "demo.py", "content": "\"\"\"\nGLiClass Enhanced Demo with Advanced Features\n\nFeatures:\n- Task description prompts\n- Hierarchical label inputs (JSON format)\n- Few-shot examples\n- Hierarchical output structure\n- Label descriptions\n\"\"\"\n\nimport json\nfrom typing import Dict, List, Any, Union, Optional\nimport gradio as gr\nimport torch\nfrom transformers import AutoTokenizer\n\nfrom gliclass import GLiClassModel, ZeroShotClassificationPipeline\n\n# Initialize model and pipeline\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\nmodel_path = \"knowledgator/gliclass-small-v1.0\"\nmodel = GLiClassModel.from_pretrained(model_path)\ntokenizer = AutoTokenizer.from_pretrained(model_path)\n\npipeline = ZeroShotClassificationPipeline(\n model, tokenizer, \n classification_type='multi-label', \n device=device\n)\n\n# ============== Example Texts ==============\n\nTEXT_PRODUCT_REVIEW = \"\"\"\nI recently purchased the Sony WH-1000XM4 Wireless Noise-Canceling Headphones from Amazon and I must say, I'm thoroughly impressed. The package arrived in New York within 2 days, thanks to Amazon Prime's expedited shipping.\n\nThe headphones themselves are remarkable. The noise-canceling feature works like a charm in the bustling city environment, and the 30-hour battery life means I don't have to charge them every day. Connecting them to my Samsung Galaxy S21 was a breeze, and the sound quality is second to none.\n\nI also appreciated the customer service from Amazon when I had a question about the warranty. They responded within an hour and provided all the information I needed.\n\nHowever, the headphones did not come with a hard case, which was listed in the product description. I contacted Amazon, and they offered a 10% discount on my next purchase as an apology.\n\nOverall, I'd give these headphones a 4.5/5 rating and highly recommend them to anyone looking for top-notch quality in both product and service.\n\"\"\"\n\nTEXT_TECH_COMPANIES = \"\"\"\nApple Inc. is an American multinational technology company headquartered in Cupertino, California. Apple is the world's largest technology company by revenue, with US$394.3 billion in 2022 revenue. As of March 2023, Apple is the world's biggest company by market capitalization.\n\nMicrosoft was founded by Bill Gates and Paul Allen on April 4, 1975 to develop and sell BASIC interpreters for the Altair 8800. During his career at Microsoft, Gates held the positions of chairman, chief executive officer, president and chief software architect.\n\nApple was founded as Apple Computer Company on April 1, 1976, by Steve Wozniak, Steve Jobs (1955–2011) and Ronald Wayne to develop and sell Wozniak's Apple I personal computer.\n\"\"\"\n\nTEXT_SCIENTIFIC = \"\"\"\nSeveral studies have reported its pharmacological activities, including anti-inflammatory, antimicrobial, and antitumoral effects. \nThe effect of E-anethole was studied in the osteosarcoma MG-63 cell line, and the antiproliferative activity was evaluated by an MTT assay. \nIt showed a GI50 value of 60.25 ΞΌM with apoptosis induction through the mitochondrial-mediated pathway. Additionally, it induced cell cycle arrest at the G0/G1 phase, up-regulated the expression of p53, caspase-3, and caspase-9, and down-regulated Bcl-xL expression.\n\"\"\"\n\nTEXT_RESTAURANT_REVIEW = \"\"\"\nWe visited La Maison last Friday for our anniversary dinner. The ambiance was absolutely stunning - dim lighting, soft jazz music, and elegant table settings. Our waiter, Marcus, was incredibly attentive without being intrusive.\n\nFor appetizers, we had the truffle bruschetta and the soup of the day. Both were divine! The main courses - filet mignon for me and lobster risotto for my wife - were cooked to perfection. \n\nThe only downside was the wait time for our desserts, which took about 25 minutes. However, the chocolate soufflΓ© was worth the wait!\n\nPrice was on the higher side ($180 for two), but the quality justified the cost. Will definitely return!\n\"\"\"\n\nTEXT_NEWS_POLITICS = \"\"\"\nThe Senate passed a landmark bipartisan infrastructure bill late Thursday night, allocating $1.2 trillion for roads, bridges, broadband internet, and clean energy initiatives. The vote was 69-30, with 19 Republican senators joining all Democrats in support.\n\nPresident Biden called the passage \"a historic investment in America's future\" and urged the House to act quickly. However, progressive Democrats have signaled they won't vote for the infrastructure bill unless it's paired with a larger social spending package.\n\nSenate Minority Leader criticized portions of the bill related to climate spending, calling them \"unnecessary green new deal provisions,\" while environmental groups praised the clean energy investments as \"a step in the right direction, but not nearly enough.\"\n\"\"\"\n\nTEXT_SPORTS = \"\"\"\nIn a thrilling overtime finish, the Lakers defeated the Celtics 118-112 in Game 7 of the NBA Finals. LeBron James delivered a historic performance with 42 points, 16 rebounds, and 10 assists, securing his fifth championship ring and fourth Finals MVP award.\n\nThe game was tied at 102 with 30 seconds remaining in regulation when Marcus Smart hit a contested three-pointer. However, James answered with a driving layup at the buzzer to force overtime.\n\nIn the extra period, the Lakers outscored Boston 16-10, with Anthony Davis contributing two crucial blocks in the final minute. \"This is what you dream about as a kid,\" James said in the post-game interview. \"Playing against the Celtics, Game 7, everything on the line.\"\n\"\"\"\n\nTEXT_MOVIE_REVIEW = \"\"\"\nChristopher Nolan's \"Oppenheimer\" is a masterwork of biographical cinema that demands to be seen on the largest screen possible. Cillian Murphy delivers a career-defining performance as J. Robert Oppenheimer, capturing both the brilliance and moral anguish of the father of the atomic bomb.\n\nThe film's nonlinear structure, weaving between the Manhattan Project, the 1954 security hearing, and the 1959 Lewis Strauss confirmation hearing, could have been confusing. Instead, Nolan crafts a compelling narrative that builds to a devastating emotional climax.\n\nAt three hours, some viewers may find the pacing challenging, particularly in the courtroom sequences. However, the technical achievements - Ludwig GΓΆransson's haunting score, Hoyte van Hoytema's IMAX cinematography - make this an unmissable theatrical experience. Rating: 9/10\n\"\"\"\n\nTEXT_TECH_STARTUP = \"\"\"\nSan Francisco-based AI startup Anthropic announced today it has raised $450 million in Series C funding, valuing the company at $5 billion. The round was led by Spark Capital, with participation from Google and existing investors.\n\nFounded in 2021 by former OpenAI researchers Dario and Daniela Amodei, Anthropic has positioned itself as a leader in AI safety research. The company's Claude assistant has gained significant market share in the enterprise segment.\n\n\"This funding will accelerate our research into interpretable and steerable AI systems,\" said CEO Dario Amodei. \"We believe safety and capability go hand in hand.\" The company plans to double its research team and expand internationally, with offices planned in London and Tokyo.\n\"\"\"\n\nTEXT_HEALTH_WELLNESS = \"\"\"\nA new study published in the Journal of the American Medical Association suggests that intermittent fasting may offer significant benefits beyond weight loss. Researchers followed 500 participants over two years and found improvements in cardiovascular health markers, insulin sensitivity, and cognitive function.\n\nParticipants who followed a 16:8 fasting protocol (eating within an 8-hour window) showed a 15% reduction in LDL cholesterol and a 20% improvement in fasting glucose levels compared to the control group.\n\nHowever, experts caution that intermittent fasting isn't suitable for everyone. \"Pregnant women, people with a history of eating disorders, and those with certain medical conditions should consult their doctor first,\" said Dr. Sarah Chen, the study's lead author. \"It's not a magic solution, but for many people, it can be a sustainable approach to improving metabolic health.\"\n\"\"\"\n\nTEXT_TRAVEL = \"\"\"\nHidden among the limestone karsts of Ha Long Bay, Cat Ba Island offers travelers an authentic Vietnamese experience away from the tourist crowds. We spent five days exploring this gem and discovered why it's becoming a favorite among backpackers and adventure seekers.\n\nThe island's national park features challenging hikes through tropical rainforest, with the trek to the peak of Ngu Lam offering panoramic views of the bay. We also kayaked through hidden lagoons and explored caves that few tourists ever see.\n\nAccommodation ranges from basic hostels ($8/night) to comfortable eco-resorts ($60/night). The seafood is incredibly fresh - we had the best grilled squid of our lives at a family-run restaurant in Cat Ba Town for just $5. Pro tip: rent a motorbike to explore the quieter beaches on the island's east side.\n\"\"\"\n\nTEXT_COOKING_RECIPE = \"\"\"\nThis Thai green curry comes together in just 30 minutes and tastes better than takeout. The secret is making your own curry paste - it takes an extra 10 minutes but the flavor difference is remarkable.\n\nFor the paste, blend together: 10 green chilies, 4 garlic cloves, 2 shallots, 1 stalk lemongrass, 1 inch galangal, handful of cilantro stems, 1 tsp cumin, 1 tsp coriander, zest of 1 lime, and 2 tbsp fish sauce. \n\nHeat coconut oil in a wok, fry the paste for 2 minutes until fragrant. Add chicken (or tofu), cook until browned. Pour in coconut milk, add bamboo shoots, Thai eggplant, and basil. Simmer for 15 minutes. Season with palm sugar and more fish sauce to taste.\n\nServe over jasmine rice with extra chilies on the side. This recipe serves 4 and can be made ahead - the flavors actually improve overnight.\n\"\"\"\n\nTEXT_FINANCIAL_ADVICE = \"\"\"\nWith inflation running at 4.2% and the Fed signaling more rate hikes, many investors are wondering how to position their portfolios. Here's what our analysis suggests for Q4 2024.\n\nFixed income is finally attractive again. With 10-year Treasury yields above 4.5%, bonds offer meaningful real returns for the first time in years. We recommend increasing allocation to investment-grade corporate bonds and TIPS for inflation protection.\n\nFor equities, we're cautiously optimistic on value stocks, particularly in the energy and financial sectors. Tech valuations remain stretched despite recent pullbacks. International developed markets, especially Japan and Europe, offer better risk-reward at current levels.\n\nRemember: past performance doesn't guarantee future results. This is general information, not personalized advice. Consult a financial advisor before making investment decisions.\n\"\"\"\n\nTEXT_ENVIRONMENTAL = \"\"\"\nThe Great Barrier Reef experienced its sixth mass bleaching event in a decade this summer, with aerial surveys showing 91% of reefs affected. Scientists warn that without dramatic action on climate change, the world's largest coral ecosystem may not survive beyond 2050.\n\n\"We're witnessing the collapse of one of Earth's most biodiverse ecosystems in real time,\" said Dr. Terry Hughes of James Cook University. Water temperatures reached 2Β°C above the February average, causing corals to expel the symbiotic algae that give them color and nutrients.\n\nSome researchers are experimenting with heat-resistant coral varieties and cloud-brightening technology to shade reefs. However, most scientists agree these are stopgap measures. \"The only real solution is rapid decarbonization,\" Hughes said. \"Everything else is just buying time.\"\n\"\"\"\n\nTEXT_EDUCATION = \"\"\"\nThe debate over standardized testing in American schools has intensified following a new report showing significant post-pandemic learning gaps. The National Assessment of Educational Progress found that fourth-grade math scores dropped to levels not seen since 2005.\n\nProponents of testing argue that standardized assessments are essential for identifying struggling students and holding schools accountable. \"Without data, we're flying blind,\" said Education Secretary Miguel Cardona. \"Tests help us direct resources where they're needed most.\"\n\nCritics counter that high-stakes testing narrows the curriculum and increases student stress without improving outcomes. \"We're testing kids more than ever, but educational outcomes aren't improving,\" said education researcher Dr. Pasi Sahlberg. \"Countries like Finland, which use minimal standardized testing, consistently outperform the US.\"\n\"\"\"\n\nTEXT_FASHION = \"\"\"\nMilan Fashion Week wrapped up yesterday with several surprising trends that will likely dominate fall/winter 2025. After years of quiet luxury and minimalism, designers are embracing bold maximalism - think dramatic volumes, clashing prints, and unapologetic color.\n\nPrada's collection featured oversized coats with exaggerated shoulders paired with flowing silk pants, while Gucci returned to its pattern-mixing roots under new creative direction. Versace went full baroque with gold-embroidered gowns that would feel at home in a Renaissance painting.\n\nSustainability remained a talking point, with Stella McCartney showcasing a collection made entirely from recycled ocean plastic. However, critics noted that the industry still has far to go. \"One sustainable collection doesn't offset the environmental impact of fast fashion,\" noted fashion journalist Vanessa Friedman. \"The industry needs systemic change, not just good PR.\"\n\"\"\"\n\nTEXT_LEGAL_CASE = \"\"\"\nThe Supreme Court agreed Monday to hear a case that could reshape the boundaries of free speech on social media platforms. The case, NetChoice v. Paxton, challenges Texas and Florida laws that prohibit large social media companies from removing certain political content.\n\nTech companies argue that the First Amendment protects their right to moderate content as they see fit, similar to how newspapers decide what to publish. \"Forcing platforms to host speech they find objectionable is compelled speech, which the Constitution forbids,\" said NetChoice counsel Paul Clement.\n\nTexas and Florida counter that social media platforms function as common carriers or public utilities and should be subject to similar non-discrimination requirements. \"These companies have become the modern public square,\" said Texas Attorney General Ken Paxton. \"They shouldn't be able to silence voices based on political viewpoint.\"\n\"\"\"\n\nTEXT_GAMING = \"\"\"\nAfter three years in development hell, \"Hollow Eclipse\" has finally launched - and it's everything fans hoped for. This action RPG from indie studio Moonlight Games delivers a haunting 40-hour adventure that rivals titles from studios with ten times the budget.\n\nThe combat system strikes a perfect balance between accessibility and depth. Basic attacks and dodges are simple to execute, but mastering the \"shadow merge\" mechanic - which lets you temporarily possess enemies - adds layers of strategy. Boss fights are challenging without feeling unfair, though the final boss may take even experienced players dozens of attempts.\n\nWhere the game truly shines is its atmosphere. The decaying gothic city of Velmoor is rendered in stunning hand-drawn art, and the ambient soundtrack creates constant unease. The story tackles themes of grief and memory with surprising emotional maturity. Minor technical issues (occasional frame drops, one softlock) can't diminish this achievement. Score: 9.5/10\n\"\"\"\n\nTEXT_REAL_ESTATE = \"\"\"\nThe housing market is sending mixed signals as we enter 2025. Existing home sales fell for the third consecutive month, down 4.1% in November, yet prices continue to climb in most metropolitan areas. The median home price hit $416,000, up 3.8% year-over-year.\n\nLow inventory remains the central issue. Many homeowners are reluctant to sell because they've locked in sub-3% mortgage rates and don't want to trade up to today's 7% rates. This \"lock-in effect\" has created a severe shortage of listings, particularly in the starter home category.\n\n\"We're seeing bidding wars even in this high-rate environment because there's simply nothing to buy,\" said economist Lawrence Yun. First-time buyers are particularly squeezed, with affordability at its worst level since 1984. Some markets, including Austin and Phoenix, are showing price corrections, but coastal cities remain stubbornly expensive.\n\"\"\"\n\nTEXT_MENTAL_HEALTH = \"\"\"\nWorkplace burnout has reached epidemic proportions, with a new Gallup survey finding that 76% of employees experience burnout at least sometimes. But recognizing burnout isn't always straightforward - it often manifests differently than simple exhaustion.\n\nThe three hallmarks of burnout are: emotional exhaustion (feeling drained and unable to cope), depersonalization (becoming cynical and detached from work), and reduced personal accomplishment (feeling ineffective regardless of actual performance).\n\nRecovery requires more than a vacation. \"You can't just rest your way out of burnout,\" says psychologist Dr. Christina Maslach, who pioneered burnout research. \"You need to address the root causes - usually workload, lack of control, insufficient recognition, or values conflicts.\" Strategies include setting firm boundaries, delegating tasks, and having honest conversations with managers about sustainable workloads. In severe cases, professional support from a therapist can help.\n\"\"\"\n\nTEXT_ASTRONOMY = \"\"\"\nNASA's James Webb Space Telescope has detected what may be signs of biological activity in the atmosphere of K2-18b, an exoplanet 120 light-years away. The discovery has electrified the scientific community, though researchers caution against jumping to conclusions.\n\nThe telescope's spectrometers identified dimethyl sulfide (DMS), a molecule produced almost exclusively by living organisms on Earth. Webb also confirmed the presence of methane and carbon dioxide, consistent with a water-rich atmosphere.\n\n\"This is tantalizing, but not definitive proof of life,\" said lead researcher Dr. Nikku Madhusudhan. \"DMS could potentially be produced by unknown geological processes. We need more observations.\" K2-18b is a \"Hycean\" world - a planet with a hydrogen-rich atmosphere and potentially a liquid water ocean beneath. If confirmed, this would be humanity's first detection of a potential biosignature beyond our solar system.\n\"\"\"\n\n\ndef parse_labels_input(labels_input: str) -> Union[List[str], Dict[str, Any]]:\n \"\"\"\n Parse labels input - supports both comma-separated and JSON hierarchical format.\n \n Examples:\n - \"positive, negative, neutral\" -> [\"positive\", \"negative\", \"neutral\"]\n - '{\"sentiment\": [\"positive\", \"negative\"], \"topic\": [\"food\", \"service\"]}' -> dict\n \"\"\"\n labels_input = labels_input.strip()\n \n # Try parsing as JSON first (for hierarchical labels)\n if labels_input.startswith('{'):\n try:\n return json.loads(labels_input)\n except json.JSONDecodeError as e:\n raise ValueError(f\"Invalid JSON format for hierarchical labels: {e}\")\n \n # Otherwise, treat as comma-separated flat labels\n labels = [label.strip() for label in labels_input.split(',') if label.strip()]\n return labels\n\n\ndef parse_examples_input(examples_input: str) -> Optional[List[Dict[str, Any]]]:\n \"\"\"\n Parse few-shot examples input (JSON format).\n \n Expected format:\n [\n {\"text\": \"Example text 1\", \"labels\": [\"label1\", \"label2\"]},\n {\"text\": \"Example text 2\", \"labels\": [\"label3\"]}\n ]\n \"\"\"\n if not examples_input or not examples_input.strip():\n return None\n \n try:\n examples = json.loads(examples_input.strip())\n if not isinstance(examples, list):\n raise ValueError(\"Examples must be a JSON array\")\n \n for i, ex in enumerate(examples):\n if not isinstance(ex, dict):\n raise ValueError(f\"Example {i+1} must be a JSON object\")\n if 'text' not in ex:\n raise ValueError(f\"Example {i+1} missing 'text' field\")\n if 'labels' not in ex and 'true_labels' not in ex:\n raise ValueError(f\"Example {i+1} missing 'labels' field\")\n \n return examples\n except json.JSONDecodeError as e:\n raise ValueError(f\"Invalid JSON format for examples: {e}\")\n\n\ndef format_output(\n results: Union[List[Dict], Dict], \n hierarchical: bool = False,\n output_format: str = \"visual\"\n) -> Union[Dict[str, float], str]:\n \"\"\"Format classification output for Gradio display.\"\"\"\n \n if output_format == \"json\":\n return format_as_json(results, hierarchical)\n \n if hierarchical and isinstance(results, dict):\n # Format hierarchical output as readable string\n return format_hierarchical_dict(results)\n \n if isinstance(results, list):\n return {result['label']: float(result['score']) for result in results}\n \n return results\n\n\ndef format_as_json(results: Union[List[Dict], Dict], hierarchical: bool = False) -> str:\n \"\"\"Format results as pretty-printed JSON string.\"\"\"\n if hierarchical and isinstance(results, dict):\n # Already in hierarchical dict format\n return json.dumps(results, indent=2, ensure_ascii=False)\n \n if isinstance(results, list):\n # Convert list of predictions to structured format\n output = {\n \"predictions\": [\n {\"label\": r[\"label\"], \"score\": round(r[\"score\"], 4)}\n for r in results\n ],\n \"scores\": {r[\"label\"]: round(r[\"score\"], 4) for r in results}\n }\n return json.dumps(output, indent=2, ensure_ascii=False)\n \n return json.dumps(results, indent=2, ensure_ascii=False)\n\n\ndef format_hierarchical_dict(d: Dict, indent: int = 0) -> str:\n \"\"\"Format hierarchical dict for display with visual score bars.\"\"\"\n lines = []\n prefix = \" \" * indent\n \n for key, value in d.items():\n if isinstance(value, dict):\n lines.append(f\"{prefix}**{key}**:\")\n lines.append(format_hierarchical_dict(value, indent + 1))\n else:\n score_bar = \"β–ˆ\" * int(value * 20) + \"β–‘\" * (20 - int(value * 20))\n lines.append(f\"{prefix}{key}: {score_bar} {value:.3f}\")\n \n return \"\\n\".join(lines)\n\n\ndef classification(\n text: str,\n labels_input: str,\n threshold: float,\n multi_label: bool,\n prompt: str,\n examples_input: str,\n hierarchical_output: bool,\n output_format: str = \"visual\"\n) -> Union[Dict[str, float], str]:\n \"\"\"\n Perform classification with all advanced features.\n \"\"\"\n try:\n # Parse labels (flat or hierarchical)\n labels = parse_labels_input(labels_input)\n \n # Parse few-shot examples\n examples = parse_examples_input(examples_input) if examples_input else None\n \n # Set classification type\n pipeline.pipe.classification_type = 'multi-label' if multi_label else 'single-label'\n \n # Prepare prompt\n task_prompt = prompt.strip() if prompt and prompt.strip() else None\n \n # Run classification\n results = pipeline(\n text, \n labels, \n threshold=threshold,\n examples=examples,\n prompt=task_prompt,\n return_hierarchical=hierarchical_output\n )[0] # Single text, get first result\n \n # Format output based on selected format\n if output_format == \"json\":\n return format_as_json(results, hierarchical_output)\n elif hierarchical_output:\n return format_hierarchical_dict(results)\n else:\n return {result['label']: float(result['score']) for result in results}\n \n except Exception as e:\n return f\"Error: {str(e)}\"\n\n\n# ============== Example Configurations ==============\n\nEXAMPLES = [\n # Example 1: Basic flat labels with prompt\n [\n TEXT_PRODUCT_REVIEW,\n \"product review, electronics, positive feedback, negative feedback, customer service, shipping\",\n 0.5,\n True,\n \"Classify this customer review by topic and sentiment:\",\n \"\",\n False,\n \"visual\"\n ],\n # Example 2: Hierarchical labels for restaurant review\n [\n TEXT_RESTAURANT_REVIEW,\n '''{\n \"sentiment\": [\"positive\", \"negative\", \"mixed\"],\n \"aspects\": [\"food quality\", \"service\", \"ambiance\", \"price\", \"wait time\"],\n \"recommendation\": [\"would recommend\", \"would not recommend\"]\n}''',\n 0.4,\n True,\n \"Analyze this restaurant review:\",\n \"\",\n True,\n \"visual\"\n ],\n # Example 3: News article with few-shot examples\n [\n TEXT_NEWS_POLITICS,\n \"politics, business, technology, sports, entertainment, science, health\",\n 0.5,\n True,\n \"Classify this news article by category:\",\n '''[\n {\"text\": \"The Federal Reserve raised interest rates by 0.25% today, citing persistent inflation concerns.\", \"labels\": [\"politics\", \"business\"]},\n {\"text\": \"Scientists discover high new high-temperature superconductor material that works at room temperature.\", \"labels\": [\"science\", \"technology\"]}\n]''',\n False,\n \"visual\"\n ],\n # Example 4: Scientific classification with hierarchical output\n [\n TEXT_SCIENTIFIC,\n '''{\n \"domain\": [\"biology\", \"chemistry\", \"medicine\", \"physics\"],\n \"research_type\": [\"experimental\", \"theoretical\", \"review\"],\n \"application\": [\"therapeutic\", \"diagnostic\", \"basic research\"]\n}''',\n 0.3,\n True,\n \"Classify this scientific abstract:\",\n \"\",\n True,\n \"visual\"\n ],\n # Example 5: Sports article - single label\n [\n TEXT_SPORTS,\n \"basketball, football, soccer, tennis, baseball, hockey, golf\",\n 0.5,\n False,\n \"What sport is this article about?\",\n \"\",\n False,\n \"visual\"\n ],\n # Example 6: Movie review with detailed sentiment (JSON output)\n [\n TEXT_MOVIE_REVIEW,\n '''{\n \"overall_sentiment\": [\"positive\", \"negative\", \"mixed\"],\n \"aspects_praised\": [\"acting\", \"direction\", \"cinematography\", \"music\", \"story\", \"pacing\"],\n \"aspects_criticized\": [\"acting\", \"direction\", \"cinematography\", \"music\", \"story\", \"pacing\"],\n \"recommendation\": [\"must watch\", \"worth watching\", \"skip it\"]\n}''',\n 0.35,\n True,\n \"Analyze this movie review in detail:\",\n \"\",\n True,\n \"json\"\n ],\n # Example 7: Tech startup news\n [\n TEXT_TECH_STARTUP,\n \"funding announcement, product launch, acquisition, IPO, partnership, hiring, layoffs, legal\",\n 0.4,\n True,\n \"What type of tech news is this?\",\n \"\",\n False,\n \"visual\"\n ],\n # Example 8: Health article with hierarchical categories\n [\n TEXT_HEALTH_WELLNESS,\n '''{\n \"topic\": [\"nutrition\", \"exercise\", \"mental health\", \"sleep\", \"medical research\"],\n \"content_type\": [\"research findings\", \"practical advice\", \"expert opinion\", \"warning\"],\n \"audience\": [\"general public\", \"healthcare professionals\", \"patients\"]\n}''',\n 0.4,\n True,\n \"Categorize this health article:\",\n \"\",\n True,\n \"visual\"\n ],\n # Example 9: Travel content (JSON output)\n [\n TEXT_TRAVEL,\n \"destination guide, hotel review, restaurant review, adventure travel, budget travel, luxury travel, travel tips\",\n 0.4,\n True,\n \"What type of travel content is this?\",\n \"\",\n False,\n \"json\"\n ],\n # Example 10: Recipe classification\n [\n TEXT_COOKING_RECIPE,\n '''{\n \"cuisine\": [\"Thai\", \"Italian\", \"Mexican\", \"Indian\", \"Chinese\", \"Japanese\", \"French\", \"American\"],\n \"difficulty\": [\"easy\", \"medium\", \"hard\"],\n \"meal_type\": [\"breakfast\", \"lunch\", \"dinner\", \"dessert\", \"snack\"],\n \"dietary\": [\"vegetarian friendly\", \"vegan friendly\", \"gluten free\", \"dairy free\", \"contains meat\"]\n}''',\n 0.35,\n True,\n \"Classify this recipe:\",\n \"\",\n True,\n \"visual\"\n ],\n # Example 11: Financial content with examples\n [\n TEXT_FINANCIAL_ADVICE,\n \"investment advice, market analysis, personal finance, retirement planning, tax advice, economic news\",\n 0.4,\n True,\n \"Categorize this financial content:\",\n '''[\n {\"text\": \"Here are 5 ways to maximize your 401k contributions before year end.\", \"labels\": [\"personal finance\", \"retirement planning\", \"tax advice\"]},\n {\"text\": \"The S&P 500 rose 2% today following strong jobs report.\", \"labels\": [\"market analysis\", \"economic news\"]}\n]''',\n False,\n \"visual\"\n ],\n # Example 12: Environmental news (JSON output)\n [\n TEXT_ENVIRONMENTAL,\n '''{\n \"topic\": [\"climate change\", \"biodiversity\", \"pollution\", \"conservation\", \"renewable energy\"],\n \"tone\": [\"alarming\", \"hopeful\", \"neutral\", \"urgent\"],\n \"focus\": [\"problem description\", \"solutions\", \"policy\", \"research findings\"]\n}''',\n 0.35,\n True,\n \"Analyze this environmental article:\",\n \"\",\n True,\n \"json\"\n ],\n # Example 13: Education debate\n [\n TEXT_EDUCATION,\n \"education policy, standardized testing, curriculum, teacher issues, student welfare, technology in education, higher education\",\n 0.4,\n True,\n \"What education topics does this article cover?\",\n \"\",\n False,\n \"visual\"\n ],\n # Example 14: Fashion news with hierarchy\n [\n TEXT_FASHION,\n '''{\n \"content_type\": [\"trend report\", \"designer profile\", \"collection review\", \"industry news\", \"sustainability\"],\n \"season\": [\"spring/summer\", \"fall/winter\"],\n \"market_segment\": [\"luxury\", \"fast fashion\", \"sustainable fashion\", \"streetwear\"]\n}''',\n 0.4,\n True,\n \"Classify this fashion article:\",\n \"\",\n True,\n \"visual\"\n ],\n # Example 15: Legal case (JSON output)\n [\n TEXT_LEGAL_CASE,\n \"constitutional law, criminal law, civil rights, corporate law, intellectual property, free speech, privacy\",\n 0.4,\n True,\n \"What areas of law does this case involve?\",\n \"\",\n False,\n \"json\"\n ],\n # Example 16: Gaming review with detailed analysis\n [\n TEXT_GAMING,\n '''{\n \"genre\": [\"action\", \"RPG\", \"adventure\", \"puzzle\", \"strategy\", \"simulation\", \"sports\"],\n \"platform_feel\": [\"indie\", \"AAA\", \"mid-tier\"],\n \"strengths\": [\"gameplay\", \"story\", \"graphics\", \"music\", \"replayability\"],\n \"weaknesses\": [\"bugs\", \"difficulty\", \"length\", \"graphics\", \"story\"],\n \"recommendation\": [\"must play\", \"worth playing\", \"wait for sale\", \"skip\"]\n}''',\n 0.35,\n True,\n \"Analyze this game review:\",\n \"\",\n True,\n \"visual\"\n ],\n # Example 17: Real estate market analysis\n [\n TEXT_REAL_ESTATE,\n \"market analysis, buying advice, selling advice, investment, rental market, mortgage rates, housing policy\",\n 0.4,\n True,\n \"What real estate topics are covered?\",\n \"\",\n False,\n \"visual\"\n ],\n # Example 18: Mental health with few-shot (JSON output)\n [\n TEXT_MENTAL_HEALTH,\n '''{\n \"topic\": [\"burnout\", \"anxiety\", \"depression\", \"stress management\", \"work-life balance\"],\n \"content_type\": [\"educational\", \"self-help advice\", \"research summary\", \"personal story\"],\n \"actionability\": [\"provides concrete steps\", \"general awareness\", \"seeks professional help\"]\n}''',\n 0.35,\n True,\n \"Categorize this mental health content:\",\n '''[\n {\"text\": \"Feeling overwhelmed? Try the 5-4-3-2-1 grounding technique: notice 5 things you see, 4 you hear...\", \"labels\": [\"topic.anxiety\", \"topic.stress management\", \"content_type.self-help advice\", \"actionability.provides concrete steps\"]},\n {\"text\": \"A new study links social media use exceeding 3 hours daily with increased rates of depression in teens.\", \"labels\": [\"topic.depression\", \"content_type.research summary\", \"actionability.general awareness\"]}\n]''',\n True,\n \"json\"\n ],\n # Example 19: Astronomy discovery\n [\n TEXT_ASTRONOMY,\n \"exoplanets, astrobiology, cosmology, solar system, space exploration, telescopes, astrophysics\",\n 0.4,\n True,\n \"What astronomy topics are discussed?\",\n \"\",\n False,\n \"visual\"\n ],\n # Example 20: Tech companies - single label\n [\n TEXT_TECH_COMPANIES,\n \"company profile, product announcement, financial report, industry analysis, biography, opinion piece\",\n 0.5,\n False,\n \"What is the primary type of this article?\",\n \"\",\n False,\n \"visual\"\n ],\n]\n\n\n# ============== Gradio Interface ==============\n\nwith gr.Blocks(\n title=\"GLiClass Advanced Demo\",\n theme=gr.themes.Soft(\n primary_hue=\"blue\",\n secondary_hue=\"slate\",\n )\n) as demo:\n \n gr.Markdown(\"\"\"\n # 🏷️ GLiClass Advanced Zero-Shot Classification\n \n Enhanced demo featuring **prompts**, **hierarchical labels**, **few-shot examples**, and **structured outputs**.\n \"\"\")\n \n with gr.Accordion(\"πŸ“– How to Use This Demo\", open=False):\n gr.Markdown(\"\"\"\n ## Features Overview\n \n ### 1. Task Description Prompts\n Add a natural language description of the classification task to guide the model.\n \n **Example:** `\"Classify this customer review by sentiment and topic:\"`\n \n ---\n \n ### 2. Hierarchical Labels (JSON Format)\n Structure your labels in categories for organized classification:\n \n ```json\n {\n \"sentiment\": [\"positive\", \"negative\", \"neutral\"],\n \"topic\": [\"product\", \"service\", \"shipping\"],\n \"urgency\": [\"high\", \"medium\", \"low\"]\n }\n ```\n \n Or use simple comma-separated labels: `positive, negative, neutral`\n \n ---\n \n ### 3. Few-Shot Examples\n Provide examples to guide the model's understanding:\n \n ```json\n [\n {\"text\": \"Great product, love it!\", \"labels\": [\"positive\", \"product\"]},\n {\"text\": \"Shipping was delayed by 2 weeks\", \"labels\": [\"negative\", \"shipping\"]}\n ]\n ```\n \n ---\n \n ### 4. Hierarchical Output\n When enabled with hierarchical labels, returns structured scores matching your input format.\n \"\"\")\n \n with gr.Accordion(\"πŸ’» Code Example\", open=False):\n gr.Code(\n '''from gliclass import GLiClassModel, ZeroShotClassificationPipeline\nfrom transformers import AutoTokenizer\n\nmodel = GLiClassModel.from_pretrained(\"knowledgator/gliclass-small-v1\")\ntokenizer = AutoTokenizer.from_pretrained(\"knowledgator/gliclass-small-v1\")\n\npipeline = ZeroShotClassificationPipeline(model, tokenizer, classification_type='multi-label', device='cuda:0')\n\n# Basic usage\ntext = \"The product quality is amazing but delivery was slow\"\nlabels = [\"positive\", \"negative\", \"product\", \"shipping\"]\nresults = pipeline(text, labels, threshold=0.5)[0]\n\n# With hierarchical labels\nhierarchical_labels = {\n \"sentiment\": [\"positive\", \"negative\", \"neutral\"],\n \"topic\": [\"product\", \"service\", \"shipping\"]\n}\n\nresults = pipeline(\n text, \n hierarchical_labels,\n prompt=\"Classify this review:\",\n return_hierarchical=True\n)[0]\n\n# With few-shot examples\nexamples = [\n {\"text\": \"Love this item!\", \"labels\": [\"sentiment.positive\", \"topic.product\"]},\n {\"text\": \"Terrible customer support\", \"labels\": [\"sentiment.negative\", \"topic.service\"]}\n]\n\nresults = pipeline(\n text,\n hierarchical_labels, \n examples=examples,\n prompt=\"Classify customer feedback:\"\n)[0]\n''',\n language=\"python\",\n )\n \n with gr.Row():\n with gr.Column(scale=2):\n input_text = gr.Textbox(\n value=EXAMPLES[0][0],\n label=\"πŸ“ Text Input\",\n placeholder=\"Enter the text you want to classify...\",\n lines=8\n )\n \n prompt_input = gr.Textbox(\n value=EXAMPLES[0][4],\n label=\"πŸ’‘ Task Description Prompt (Optional)\",\n placeholder=\"E.g., 'Classify this customer review by sentiment and topic:'\",\n lines=1\n )\n \n with gr.Column(scale=1):\n labels_input = gr.Textbox(\n value=EXAMPLES[0][1],\n label=\"🏷️ Labels (comma-separated or JSON)\",\n placeholder='positive, negative\\n\\nOR\\n\\n{\"category\": [\"label1\", \"label2\"]}',\n lines=6\n )\n \n with gr.Row():\n threshold = gr.Slider(\n 0, 1,\n value=0.5,\n step=0.01,\n label=\"Threshold\",\n info=\"Confidence threshold for predictions\"\n )\n \n with gr.Row():\n multi_label = gr.Checkbox(\n value=True,\n label=\"Multi-label\",\n info=\"Allow multiple labels per text\"\n )\n hierarchical_output = gr.Checkbox(\n value=False,\n label=\"Hierarchical Output\",\n info=\"Return structured output matching label hierarchy\"\n )\n \n with gr.Row():\n output_format = gr.Radio(\n choices=[\"visual\", \"json\"],\n value=\"visual\",\n label=\"Output Format\",\n info=\"Visual: charts/bars | JSON: raw data\"\n )\n \n with gr.Accordion(\"🎯 Few-Shot Examples (Optional)\", open=False):\n examples_input = gr.Textbox(\n value=\"\",\n label=\"Examples (JSON format)\",\n placeholder='''[\n {\"text\": \"Example text 1\", \"labels\": [\"label1\", \"label2\"]},\n {\"text\": \"Example text 2\", \"labels\": [\"label3\"]}\n]''',\n lines=5\n )\n gr.Markdown(\"\"\"\n *Provide labeled examples to guide the model. Each example needs a `text` field and a `labels` array.*\n \"\"\")\n \n submit_btn = gr.Button(\"πŸš€ Classify\", variant=\"primary\", size=\"lg\")\n \n output = gr.Label(label=\"πŸ“Š Classification Results\")\n output_text = gr.Textbox(\n label=\"πŸ“Š Hierarchical Results\", \n visible=False, \n lines=10\n )\n output_json = gr.Code(\n label=\"πŸ“Š JSON Output\",\n language=\"json\",\n visible=False,\n lines=15\n )\n \n # Dynamic output visibility based on format and hierarchical toggle\n def update_output_visibility(hierarchical: bool, fmt: str):\n if fmt == \"json\":\n return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)\n elif hierarchical:\n return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)\n else:\n return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)\n \n hierarchical_output.change(\n fn=update_output_visibility,\n inputs=[hierarchical_output, output_format],\n outputs=[output, output_text, output_json]\n )\n \n output_format.change(\n fn=update_output_visibility,\n inputs=[hierarchical_output, output_format],\n outputs=[output, output_text, output_json]\n )\n \n # Classification function wrapper for different outputs\n def classify_wrapper(text, labels, threshold, multi_label, prompt, examples, hierarchical, fmt):\n result = classification(text, labels, threshold, multi_label, prompt, examples, hierarchical, fmt)\n \n if fmt == \"json\":\n return None, None, result\n elif hierarchical or isinstance(result, str):\n return None, result, None\n else:\n return result, None, None\n \n # Event handlers\n submit_btn.click(\n fn=classify_wrapper,\n inputs=[input_text, labels_input, threshold, multi_label, prompt_input, examples_input, hierarchical_output, output_format],\n outputs=[output, output_text, output_json]\n )\n \n input_text.submit(\n fn=classify_wrapper,\n inputs=[input_text, labels_input, threshold, multi_label, prompt_input, examples_input, hierarchical_output, output_format],\n outputs=[output, output_text, output_json]\n )\n \n gr.Markdown(\"### πŸ“š Example Configurations\")\n \n gr.Examples(\n examples=EXAMPLES,\n inputs=[input_text, labels_input, threshold, multi_label, prompt_input, examples_input, hierarchical_output, output_format],\n outputs=[output, output_text, output_json],\n fn=classify_wrapper,\n cache_examples=False,\n examples_per_page=5\n )\n \n gr.Markdown(\"\"\"\n ---\n \n ### πŸ”§ Tips for Best Results\n \n | Feature | Best Practice |\n |---------|---------------|\n | **Prompts** | Be specific about the task, e.g., \"Classify by sentiment:\" vs \"Analyze:\" |\n | **Labels** | Use descriptive labels; \"customer service issue\" > \"service\" |\n | **Hierarchical** | Group related labels under categories for organized results |\n | **Examples** | 2-3 diverse examples often improve accuracy significantly |\n | **Threshold** | Start at 0.5, lower for more predictions, raise for higher precision |\n \"\"\")\n\n\nif __name__ == \"__main__\":\n demo.queue()\n demo.launch(debug=True, share=True)" }, { "path": "gliclass/__init__.py", "content": "from .model import GLiClassModel, GLiClassBiEncoder, GLiClassUniEncoder, GLiClassEncoderDecoderCLS\nfrom .config import GLiClassModelConfig\nfrom .pipeline import (\n ZeroShotClassificationPipeline,\n BiEncoderZeroShotClassificationPipeline,\n ZeroShotClassificationWithChunkingPipeline,\n)\n\n__version__ = \"0.1.19\"\n\n# Serve module (optional import)\ntry:\n from . import serve\nexcept ImportError:\n serve = None\n" }, { "path": "gliclass/config.py", "content": "from transformers import AutoConfig\nfrom transformers.utils import logging\nfrom transformers.models.auto import CONFIG_MAPPING\nfrom transformers.configuration_utils import PretrainedConfig\n\nfrom .utils import is_module_available\n\nIS_TURBOT5 = is_module_available(\"turbot5\")\n\nif IS_TURBOT5:\n from turbot5.model.config import T5Config\nelse:\n from transformers import T5Config\n\n\nlogger = logging.get_logger(__name__)\n\n\nclass GLiClassModelConfig(PretrainedConfig):\n model_type = \"GLiClass\"\n is_composition = True\n\n def __init__(\n self,\n encoder_config=None,\n encoder_model=None,\n label_model_config=None,\n label_model_name=None,\n class_token_index=-1,\n text_token_index=-1,\n example_token_index=-1,\n ignore_index=-100,\n hidden_size=None,\n projector_hidden_act=\"gelu\",\n vocab_size=None,\n problem_type=\"single_label_classification\",\n max_num_classes=25,\n use_lstm=False,\n initializer_range=0.03,\n scorer_type=\"simple\",\n scorer_num_heads=16,\n scorer_mlp_hidden_size=1024,\n scorer_attn_dropout=0.1,\n pooling_strategy=\"first\",\n class_token_pooling=\"first\",\n focal_loss_alpha=0.5,\n focal_loss_gamma=2,\n focal_loss_reduction=None,\n logit_scale_init_value=2.6592,\n normalize_features=False,\n extract_text_features=False,\n max_labels_alloc: str = \"dynamic\",\n contrastive_loss_coef=0,\n architecture_type=\"uni-encoder\",\n prompt_first=False,\n squeeze_layers=False,\n layer_wise=False,\n encoder_layer_id=-1,\n embed_class_token=True,\n dropout=0.1,\n use_segment_embeddings=False,\n **kwargs,\n ):\n if isinstance(encoder_config, dict):\n encoder_config[\"model_type\"] = encoder_config.get(\"model_type\", \"deberta-v2\")\n if encoder_config[\"model_type\"] == \"t5\":\n encoder_config = T5Config(**encoder_config)\n elif encoder_config[\"model_type\"] in CONFIG_MAPPING:\n encoder_config = CONFIG_MAPPING[encoder_config[\"model_type\"]](**encoder_config)\n else:\n _name = encoder_model or kwargs.get(\"encoder_model_name\")\n if _name:\n encoder_config = AutoConfig.from_pretrained(_name, trust_remote_code=True)\n else:\n encoder_config = PretrainedConfig(**encoder_config)\n elif encoder_config is None:\n encoder_config = CONFIG_MAPPING[\"deberta-v2\"]()\n\n self.encoder_config = encoder_config\n self.encoder_model_name = encoder_model\n\n if label_model_name is not None:\n if isinstance(label_model_config, dict):\n label_model_config[\"model_type\"] = label_model_config.get(\"model_type\", \"deberta-v2\")\n label_model_config = CONFIG_MAPPING[label_model_config[\"model_type\"]](**label_model_config)\n elif label_model_config is None:\n label_model_config = CONFIG_MAPPING[\"deberta-v2\"]()\n\n self.label_model_config = label_model_config\n else:\n self.label_model_config = None\n self.label_model_name = label_model_name\n\n if hidden_size is None:\n self.hidden_size = self.encoder_config.hidden_size\n else:\n self.hidden_size = hidden_size\n\n if vocab_size is None:\n self.vocab_size = self.encoder_config.vocab_size\n else:\n self.vocab_size = vocab_size\n\n if class_token_index == -1:\n self.class_token_index = self.vocab_size\n else:\n self.class_token_index = class_token_index\n\n if text_token_index == -1:\n self.text_token_index = self.vocab_size + 1\n else:\n self.text_token_index = text_token_index\n\n if example_token_index == -1:\n self.example_token_index = self.vocab_size + 2\n else:\n self.example_token_index = example_token_index\n\n self.ignore_index = ignore_index\n self.projector_hidden_act = projector_hidden_act\n self.problem_type = problem_type\n self.max_num_classes = max_num_classes\n self.initializer_range = initializer_range\n self.scorer_type = scorer_type\n self.scorer_num_heads = scorer_num_heads\n self.scorer_mlp_hidden_size = scorer_mlp_hidden_size\n self.scorer_attn_dropout = scorer_attn_dropout\n self.pooling_strategy = pooling_strategy\n self.class_token_pooling = class_token_pooling\n self.use_lstm = use_lstm\n self.focal_loss_alpha = focal_loss_alpha\n self.focal_loss_gamma = focal_loss_gamma\n self.focal_loss_reduction = focal_loss_reduction\n self.contrastive_loss_coef = contrastive_loss_coef\n self.logit_scale_init_value = logit_scale_init_value\n self.normalize_features = normalize_features\n self.extract_text_features = extract_text_features\n self.max_labels_alloc = max_labels_alloc\n self.architecture_type = architecture_type\n self.prompt_first = prompt_first\n self.squeeze_layers = squeeze_layers\n self.layer_wise = layer_wise\n self.encoder_layer_id = encoder_layer_id\n self.embed_class_token = embed_class_token\n self.pad_token_id = self.encoder_config.pad_token_id\n self.dropout = dropout\n self.use_segment_embeddings = use_segment_embeddings\n super().__init__(**kwargs)\n" }, { "path": "gliclass/data_processing.py", "content": "import copy\nimport random\nfrom dataclasses import dataclass\n\nimport torch\nfrom torch.utils.data import Dataset\nfrom torch.nn.utils.rnn import pad_sequence\n\n\n@dataclass\nclass AugmentationConfig:\n \"\"\"Configuration for data augmentation.\"\"\"\n\n enabled: bool = True\n\n # Probability for each augmentation type\n random_label_removal_prob: float = 0.15\n random_label_addition_prob: float = 0.10\n random_text_addition_prob: float = 0.05\n random_add_description_prob: float = 0.25\n random_add_synonyms_prob: float = 0.1\n random_add_examples_prob: float = 0.25\n max_num_examples: int = 5\n\n\nclass DataAugmenter:\n def __init__(self, config, examples, labels, label2description=None):\n self.config = config\n self.examples = examples\n self.labels = sorted(labels)\n self.max_examples = self.config.max_num_examples\n self.label2description = label2description or {}\n\n def remove_labels(self, true_labels, all_labels):\n if len(all_labels) <= 1:\n return true_labels, all_labels\n k = random.randint(1, len(all_labels))\n all_labels = random.sample(all_labels, k=k)\n true_labels = [lbl for lbl in true_labels if lbl in all_labels]\n return true_labels, all_labels\n\n def add_random_labels(self, all_labels):\n if not self.labels:\n return all_labels\n num_add = len(all_labels) + 1\n k = random.randint(1, min(num_add, len(self.labels)))\n add_labels = random.sample(self.labels, k=k)\n all_labels.extend(add_labels)\n return all_labels\n\n def add_random_text(self, text, all_labels):\n if not self.examples:\n return text\n example = random.sample(self.examples, k=1)[0]\n curr_labels = example[\"all_labels\"]\n joint_labels = set(all_labels) & set(curr_labels)\n if len(joint_labels):\n return text\n else:\n if random.randint(0, 1):\n text = example[\"text\"] + \" \" + text\n else:\n text = text + \" \" + example[\"text\"]\n return text\n\n def add_random_synonyms(self, all_labels):\n \"\"\"Replace some labels with their synonyms if available.\"\"\"\n if not self.label2description:\n return all_labels\n\n augmented_labels = []\n for label in all_labels:\n if label in self.label2description:\n label_info = self.label2description[label]\n synonyms = label_info.get(\"synonyms\", [])\n\n if synonyms and random.random() < 0.5:\n augmented_labels.append(random.choice(synonyms))\n else:\n augmented_labels.append(label)\n else:\n augmented_labels.append(label)\n\n return augmented_labels\n\n def add_random_descriptions(self, item):\n \"\"\"Add descriptions to labels in the text or metadata.\"\"\"\n if not self.label2description or not item[\"all_labels\"]:\n return item\n\n max_labels = min(3, len(item[\"all_labels\"]))\n labels_to_describe = random.sample(item[\"all_labels\"], k=random.randint(1, max_labels))\n\n descriptions = []\n for label in labels_to_describe:\n if label in self.label2description:\n label_info = self.label2description[label]\n desc_list = label_info.get(\"descriptions\", [])\n if desc_list:\n descriptions.append(f\"{label}: {random.choice(desc_list)}\")\n\n if descriptions:\n desc_text = \" \".join(descriptions)\n if random.random() < 0.5:\n item[\"text\"] = desc_text + \" \" + item[\"text\"]\n else:\n item[\"text\"] = item[\"text\"] + \" \" + desc_text\n\n return item\n\n def add_random_examples(self, item):\n \"\"\"Add example texts with similar labels.\"\"\"\n if not item[\"all_labels\"]:\n return item\n\n candidate_examples = item.get(\"examples\", [])\n\n item_label_set = set(item[\"all_labels\"])\n\n if not candidate_examples:\n for example in self.examples:\n example_label_set = set(example[\"true_labels\"])\n example_text = example[\"text\"]\n\n overlap = item_label_set & example_label_set\n\n # Only consider examples with at least one overlapping label\n if overlap:\n candidate_examples.append({\"text\": example_text, \"labels\": list(example_label_set)})\n\n if not candidate_examples:\n return item\n\n # Sort by overlap and select top examples\n random.shuffle(candidate_examples)\n top_candidates = candidate_examples[: self.max_examples]\n\n num_examples = random.randint(1, min(2, len(top_candidates)))\n selected_examples = random.sample(top_candidates, k=num_examples)\n\n item[\"examples\"] = selected_examples\n\n return item\n\n def augment(self, item):\n if not self.config.enabled:\n return item\n\n text = copy.deepcopy(item[\"text\"])\n true_labels = copy.deepcopy(item[\"true_labels\"])\n all_labels = copy.deepcopy(item[\"all_labels\"])\n\n # Create augmented item\n aug_item = {\"text\": text, \"true_labels\": true_labels, \"all_labels\": all_labels}\n\n # Copy any additional fields\n for key in item:\n if key not in aug_item:\n aug_item[key] = copy.deepcopy(item[key])\n\n if random.random() < self.config.random_label_removal_prob:\n aug_item[\"true_labels\"], aug_item[\"all_labels\"] = self.remove_labels(\n aug_item[\"true_labels\"], aug_item[\"all_labels\"]\n )\n\n if random.random() < self.config.random_label_addition_prob:\n aug_item[\"all_labels\"] = self.add_random_labels(aug_item[\"all_labels\"])\n\n if random.random() < self.config.random_text_addition_prob:\n aug_item[\"text\"] = self.add_random_text(aug_item[\"text\"], aug_item[\"all_labels\"])\n\n if random.random() < self.config.random_add_synonyms_prob:\n aug_item[\"all_labels\"] = self.add_random_synonyms(aug_item[\"all_labels\"])\n\n if random.random() < self.config.random_add_description_prob:\n aug_item = self.add_random_descriptions(aug_item)\n\n if random.random() < self.config.random_add_examples_prob:\n aug_item = self.add_random_examples(aug_item)\n\n return aug_item\n\n\nclass GLiClassDataset(Dataset):\n def __init__(\n self,\n examples,\n tokenizer,\n augment_config,\n label2description={},\n max_length=512,\n problem_type=\"multi_label_classification\",\n architecture_type=\"uni-encoder\",\n add_description=True,\n prompt_first=False,\n get_negatives=False,\n max_labels=50,\n labels_tokenizer=None,\n shuffle_labels=True,\n ):\n self.tokenizer = tokenizer\n self.labels_tokenizer = labels_tokenizer\n self.label2description = label2description\n self.augment_config = augment_config\n self.max_length = max_length\n self._data = examples\n self.add_description = add_description\n self.problem_type = problem_type\n self.architecture_type = architecture_type\n self.prompt_first = prompt_first\n self.dataset_labels = self.collect_dataset_labels()\n self.get_negatives = get_negatives\n self.max_labels = max_labels\n self.shuffle_labels = shuffle_labels\n\n self.sep_token = \"<>\"\n self.label_token = \"<