# -*- coding: utf-8 -*-
from setuptools import setup

packages = \
['rankgen', 'rankgen.parallel']

package_data = \
{'': ['*'], 'rankgen.parallel': ['parallel_logs/*', 'parallel_schedulers/*']}

install_requires = \
['gdown>=4.5.1,<5.0.0',
 'sentencepiece>=0.1.96,<0.2.0',
 'torch>=1.12.0,<2.0.0',
 'transformers>=4.20.1,<5.0.0']

setup_kwargs = {
    'name': 'rankgen',
    'version': '0.1.1',
    'description': 'RankGen is a suite of encoder models (100M-1.2B parameters) which map prefixes and generations from any pretrained English language model to a shared vector space. RankGen can be used to rerank multiple full-length samples from an LM, and it can also be incorporated as a scoring function into beam search to significantly improve generation quality (0.85 vs 0.77 MAUVE, 75% preference according to humans annotators who are English writers).',
    'long_description': '## RankGen - Improving Text Generation with Large Ranking Models\n\n[![made-with-python](https://img.shields.io/badge/Made%20with-Python-red.svg)](#python)\n[![arxiv](https://img.shields.io/badge/arXiv-2205.09726-b31b1b.svg)](https://arxiv.org/abs/2205.09726)\n[![PyPI version rankgen](https://badge.fury.io/py/rankgen.svg)](https://pypi.python.org/pypi/rankgen/) [![License: Apache 2.0](https://img.shields.io/badge/License-Apache--2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)\n\nThis is the official repository for our preprint, [RankGen - Improving Text Generation with Large Ranking Models](https://arxiv.org/abs/2205.09726). RankGen is a 1.2 billion encoder model which maps prefixes and generations from any pretrained English language model to a shared vector space. RankGen can be used to rerank multiple full-length samples from an LM, and it can also be incorporated as a scoring function into beam search to significantly improve generation quality (0.85 vs 0.77 [MAUVE](https://arxiv.org/abs/2102.01454), 75% preference according to humans annotators who are English writers). RankGen can also be used like a dense retriever, and achieves state-of-the-art performance on [literary retrieval](https://relic.cs.umass.edu/leaderboard.html).\n\nThis repository contains human evaluation data, links to HuggingFace-compatible model checkpoints, and code to integrate RankGen in beam search on HuggingFace models. RankGen is trained by fine-tuning the T5-XL encoder using the [T5X library](https://github.com/google-research/t5x).\n\n### Updates\n\n* (July 2022) RankGen is now a [PyPI package](https://pypi.org/project/rankgen), just run `pip install rankgen` to use it!\n* (July 2022) RankGen checkpoints are now available on the HuggingFace Model Hub ([link](https://huggingface.co/kalpeshk2011))!\n\n### Model checkpoints\n\nAll RankGen checkpoints are available on the HuggingFace Model Hub - [link](https://huggingface.co/kalpeshk2011)\n\nWe recommend using `RankGen-XL-all`.\n\n| Checkpoint        | Size | Hub Model Name                    | HF Hub Link                                                      |\n|-------------------|------|-----------------------------------|------------------------------------------------------------------|\n| RankGen-base-all  | 0.1B | kalpeshk2011/rankgen-t5-base-all  | [link](https://huggingface.co/kalpeshk2011/rankgen-t5-base-all)  |\n| RankGen-large-all | 0.3B | kalpeshk2011/rankgen-t5-large-all | [link](https://huggingface.co/kalpeshk2011/rankgen-t5-large-all) |\n| RankGen-XL-all    | 1.2B | kalpeshk2011/rankgen-t5-xl-all    | [link](https://huggingface.co/kalpeshk2011/rankgen-t5-xl-all)    |\n| RankGen-XL-PG19   | 1.2B | kalpeshk2011/rankgen-t5-xl-pg19   | [link](https://huggingface.co/kalpeshk2011/rankgen-t5-xl-pg19)   |\n\n*Older versions of the checkpoints*:\n\nRankGen XL checkpoints compatible with `T5XEmbeddingGeneratorLegacy` - [here](https://drive.google.com/drive/folders/1m8ujkAqkBBWYAJISZigz1Lw4tQGbZXaY?usp=sharing)\n\nT5X JAX checkpoints (base, large, XL) - [here](https://github.com/google-research/google-research/tree/master/rankgen)\n\n### Setup\n\n**Requirements**\n\nPython 3.7+, `torch` (CUDA recommended), `transformers`\n\n**Installation**\n\n(from PyPI)\n\n```\npython3.7 -m virtualenv rankgen-venv\nsource rankgen-venv/bin/activate\npip install rankgen\n```\n\n(from source)\n\n```\npython3.7 -m virtualenv rankgen-venv\nsource rankgen-venv/bin/activate\ngit clone https://github.com/martiansideofthemoon/rankgen\ncd rankgen\npip install --editable .\n```\n\n**Data Download / Test**\n\nGet the data [here](https://drive.google.com/drive/folders/1DRG2ess7fK3apfB-6KoHb_azMuHbsIv4?usp=sharing) and place folder in root directory. Alternatively, use `gdown` as shown below,\n\n```\ngdown --folder https://drive.google.com/drive/folders/1DRG2ess7fK3apfB-6KoHb_azMuHbsIv4\n```\n\nRun the test script to make sure the RankGen checkpoint has loaded correctly,\n\n```\npython -m rankgen.test_rankgen_encoder --model_path kalpeshk2011/rankgen-t5-base-all\n\n### Expected output\n0.0009239262409127233\n0.0011521980725477804\n```\n\n### Using RankGen\n\nLoading RankGen is simple using the HuggingFace APIs, but we suggest using [`RankGenEncoder`](rankgen/rankgen_encoder.py), which is a small wrapper around the HuggingFace APIs for correctly preprocessing data and doing tokenization automatically. Please see [`rankgen/test_rankgen_encoder.py`](rankgen/test_rankgen_encoder.py) for an example of the usage or see below.\n\n```\nfrom rankgen import RankGenEncoder, RankGenGenerator\n\nrankgen_encoder = RankGenEncoder("kalpeshk2011/rankgen-t5-xl-all")\n```\n\n**Encoding text to prefix/suffix vectors**\n\n```\nprefix_vectors = rankgen_encoder.encode(["This is a prefix sentence."], vectors_type="prefix")\nsuffix_vectors = rankgen_encoder.encode(["This is a suffix sentence."], vectors_type="suffix")\n```\n\n**Generating text**\n\n```\n# use a HuggingFace compatible language model\ngenerator = RankGenGenerator(rankgen_encoder=rankgen_encoder, language_model="gpt2-medium")\n\ninputs = ["Whatever might be the nature of the tragedy it would be over with long before this, and those moving black spots away yonder to the west, that he had discerned from the bluff, were undoubtedly the departing raiders. There was nothing left for Keith to do except determine the fate of the unfortunates, and give their bodies decent burial. That any had escaped, or yet lived, was altogether unlikely, unless, perchance, women had been in the party, in which case they would have been borne away prisoners."]\n\n# Baseline nucleus sampling\nprint(generator.generate_single(inputs, top_p=0.9)[0][0])\n# Over-generate and re-rank\nprint(generator.overgenerate_rerank(inputs, top_p=0.9, num_samples=10)[0][0])\n# Beam search\nprint(generator.beam_search(inputs, top_p=0.9, num_samples=10, beam_size=2)[0][0])\n```\n\n### Running beam search with RankGen (reproducing experiments in the paper)\n\nThe main file is [`rankgen/rankgen_beam_search.py`](rankgen/rankgen_beam_search.py). To execute it,\n\n```\npython rankgen/rankgen_beam_search.py \\\n    --dataset rankgen_data/wiki.jsonl \\\n    --rankgen_encoder kalpeshk2011/rankgen-t5-xl-all \\\n    --num_tokens 20 --num_samples 10 --beam_size 2 \\\n    --output_file outputs_beam/wiki_t5_xl_beam_2_tokens_20_samples_10.jsonl\n```\n\nEvaluating using MAUVE (make sure JSONL file has several thousand generations for intuitive MAUVE scores, 7713 in our experiments),\n\n```\npython rankgen/score_multi_beam.py --dataset outputs_beam/wiki_t5_xl_beam_2_tokens_10_samples_10.jsonl\n```\n\n\n### Human evaluation data\n\nWe conducted our human evaluation on Upwork, hiring English teachers and writers. We performed blind A/B testing between RankGen and nucleus sampling. We also asked our annotators to provide a 1-3 sentence explanation. You can find all the 600 annotations across two files in [`human-eval-data`](human-eval-data). To compute the evaluation scores run,\n\n```\npython rankgen/score_ab_text.py\n```\n\n### Citation Information\nIf you use RankGen, please cite it as follows:\n```\n@article{krishna2022rankgen,\n  title={RankGen: Improving Text Generation with Large Ranking Models},\n  author={Kalpesh Krishna and Yapei Chang and John Wieting and Mohit Iyyer},\n  journal={arXiv preprint arXiv:2205.09726},\n  year={2022}\n}\n```\n',
    'author': 'Kalpesh Krishna, Yapei Chang, John Wieting, Mohit Iyyer',
    'author_email': None,
    'maintainer': None,
    'maintainer_email': None,
    'url': 'https://github.com/martiansideofthemoon/rankgen',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.7,<4.0',
}


setup(**setup_kwargs)
