Full Code of LaMP-Benchmark/LaMP for AI

main 38c9d4d87629 cached
50 files
221.9 KB
53.9k tokens
227 symbols
1 requests
Download .txt
Showing preview only (236K chars total). Download the full file or copy to clipboard to get everything.
Repository: LaMP-Benchmark/LaMP
Branch: main
Commit: 38c9d4d87629
Files: 50
Total size: 221.9 KB

Directory structure:
gitextract_ofn4mhve/

├── CC-BY-NC-SA-4.0.txt
├── LaMP/
│   ├── data/
│   │   └── datasets.py
│   ├── evaluate_llm.py
│   ├── metrics/
│   │   ├── classification_metrics.py
│   │   └── generation_metrics.py
│   ├── profile_item_utilization_scorer.py
│   ├── prompts/
│   │   ├── contriever_retriever.py
│   │   ├── prompts.py
│   │   └── utils.py
│   ├── rank_profiles.py
│   ├── requirements.txt
│   ├── retriever_utilization_scorer.py
│   ├── train_llm.py
│   └── utils/
│       └── merge_with_rank.py
├── PEFT/
│   ├── data/
│   │   └── datasets.py
│   ├── evaluate_llm.py
│   ├── requirements.txt
│   └── train_peft.py
├── README.md
├── ROPG/
│   ├── data/
│   │   ├── collators.py
│   │   └── datasets.py
│   ├── models/
│   │   ├── optim.py
│   │   └── retriever.py
│   ├── prompts/
│   │   ├── contriever_retriever.py
│   │   ├── prompts.py
│   │   └── utils.py
│   ├── requirements.txt
│   ├── train_kd.py
│   ├── train_rl.py
│   ├── trainers/
│   │   └── trainer.py
│   └── utils/
│       ├── distributed.py
│       ├── log.py
│       └── util.py
├── RSPG/
│   ├── data/
│   │   ├── collators.py
│   │   └── dataset.py
│   ├── metrics/
│   │   └── evaluation.py
│   ├── modeling/
│   │   ├── __init__.py
│   │   ├── modeling.py
│   │   ├── optim.py
│   │   └── utils.py
│   ├── requirements.txt
│   ├── rspg.py
│   └── utils/
│       ├── __init__.py
│       ├── create_data.py
│       ├── distributed.py
│       └── log.py
├── data/
│   └── avocado/
│       └── create_avocado_dataset.py
└── eval/
    ├── eval_all.py
    ├── eval_task.py
    └── evaluation.py

================================================
FILE CONTENTS
================================================

================================================
FILE: CC-BY-NC-SA-4.0.txt
================================================
Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International

Creative Commons Corporation ("Creative Commons") is not a law firm and
does not provide legal services or legal advice. Distribution of
Creative Commons public licenses does not create a lawyer-client or
other relationship. Creative Commons makes its licenses and related
information available on an "as-is" basis. Creative Commons gives no
warranties regarding its licenses, any material licensed under their
terms and conditions, or any related information. Creative Commons
disclaims all liability for damages resulting from their use to the
fullest extent possible.

Using Creative Commons Public Licenses

Creative Commons public licenses provide a standard set of terms and
conditions that creators and other rights holders may use to share
original works of authorship and other material subject to copyright and
certain other rights specified in the public license below. The
following considerations are for informational purposes only, are not
exhaustive, and do not form part of our licenses.

Considerations for licensors: Our public licenses are intended for use
by those authorized to give the public permission to use material in
ways otherwise restricted by copyright and certain other rights. Our
licenses are irrevocable. Licensors should read and understand the terms
and conditions of the license they choose before applying it. Licensors
should also secure all rights necessary before applying our licenses so
that the public can reuse the material as expected. Licensors should
clearly mark any material not subject to the license. This includes
other CC-licensed material, or material used under an exception or
limitation to copyright. More considerations for licensors :
wiki.creativecommons.org/Considerations_for_licensors

Considerations for the public: By using one of our public licenses, a
licensor grants the public permission to use the licensed material under
specified terms and conditions. If the licensor's permission is not
necessary for any reason–for example, because of any applicable
exception or limitation to copyright–then that use is not regulated by
the license. Our licenses grant only permissions under copyright and
certain other rights that a licensor has authority to grant. Use of the
licensed material may still be restricted for other reasons, including
because others have copyright or other rights in the material. A
licensor may make special requests, such as asking that all changes be
marked or described. Although not required by our licenses, you are
encouraged to respect those requests where reasonable. More
considerations for the public :
wiki.creativecommons.org/Considerations_for_licensees

Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
Public License

By exercising the Licensed Rights (defined below), You accept and agree
to be bound by the terms and conditions of this Creative Commons
Attribution-NonCommercial-ShareAlike 4.0 International Public License
("Public License"). To the extent this Public License may be interpreted
as a contract, You are granted the Licensed Rights in consideration of
Your acceptance of these terms and conditions, and the Licensor grants
You such rights in consideration of benefits the Licensor receives from
making the Licensed Material available under these terms and conditions.

Section 1 – Definitions.

-   a. Adapted Material means material subject to Copyright and Similar
    Rights that is derived from or based upon the Licensed Material and
    in which the Licensed Material is translated, altered, arranged,
    transformed, or otherwise modified in a manner requiring permission
    under the Copyright and Similar Rights held by the Licensor. For
    purposes of this Public License, where the Licensed Material is a
    musical work, performance, or sound recording, Adapted Material is
    always produced where the Licensed Material is synched in timed
    relation with a moving image.
-   b. Adapter's License means the license You apply to Your Copyright
    and Similar Rights in Your contributions to Adapted Material in
    accordance with the terms and conditions of this Public License.
-   c. BY-NC-SA Compatible License means a license listed at
    creativecommons.org/compatiblelicenses, approved by Creative Commons
    as essentially the equivalent of this Public License.
-   d. Copyright and Similar Rights means copyright and/or similar
    rights closely related to copyright including, without limitation,
    performance, broadcast, sound recording, and Sui Generis Database
    Rights, without regard to how the rights are labeled or categorized.
    For purposes of this Public License, the rights specified in Section
    2(b)(1)-(2) are not Copyright and Similar Rights.
-   e. Effective Technological Measures means those measures that, in
    the absence of proper authority, may not be circumvented under laws
    fulfilling obligations under Article 11 of the WIPO Copyright Treaty
    adopted on December 20, 1996, and/or similar international
    agreements.
-   f. Exceptions and Limitations means fair use, fair dealing, and/or
    any other exception or limitation to Copyright and Similar Rights
    that applies to Your use of the Licensed Material.
-   g. License Elements means the license attributes listed in the name
    of a Creative Commons Public License. The License Elements of this
    Public License are Attribution, NonCommercial, and ShareAlike.
-   h. Licensed Material means the artistic or literary work, database,
    or other material to which the Licensor applied this Public License.
-   i. Licensed Rights means the rights granted to You subject to the
    terms and conditions of this Public License, which are limited to
    all Copyright and Similar Rights that apply to Your use of the
    Licensed Material and that the Licensor has authority to license.
-   j. Licensor means the individual(s) or entity(ies) granting rights
    under this Public License.
-   k. NonCommercial means not primarily intended for or directed
    towards commercial advantage or monetary compensation. For purposes
    of this Public License, the exchange of the Licensed Material for
    other material subject to Copyright and Similar Rights by digital
    file-sharing or similar means is NonCommercial provided there is no
    payment of monetary compensation in connection with the exchange.
-   l. Share means to provide material to the public by any means or
    process that requires permission under the Licensed Rights, such as
    reproduction, public display, public performance, distribution,
    dissemination, communication, or importation, and to make material
    available to the public including in ways that members of the public
    may access the material from a place and at a time individually
    chosen by them.
-   m. Sui Generis Database Rights means rights other than copyright
    resulting from Directive 96/9/EC of the European Parliament and of
    the Council of 11 March 1996 on the legal protection of databases,
    as amended and/or succeeded, as well as other essentially equivalent
    rights anywhere in the world.
-   n. You means the individual or entity exercising the Licensed Rights
    under this Public License. Your has a corresponding meaning.

Section 2 – Scope.

-   a. License grant.
    -   1. Subject to the terms and conditions of this Public License,
        the Licensor hereby grants You a worldwide, royalty-free,
        non-sublicensable, non-exclusive, irrevocable license to
        exercise the Licensed Rights in the Licensed Material to:
        -   A. reproduce and Share the Licensed Material, in whole or in
            part, for NonCommercial purposes only; and
        -   B. produce, reproduce, and Share Adapted Material for
            NonCommercial purposes only.
    -   2. Exceptions and Limitations. For the avoidance of doubt, where
        Exceptions and Limitations apply to Your use, this Public
        License does not apply, and You do not need to comply with its
        terms and conditions.
    -   3. Term. The term of this Public License is specified in Section
        6(a).
    -   4. Media and formats; technical modifications allowed. The
        Licensor authorizes You to exercise the Licensed Rights in all
        media and formats whether now known or hereafter created, and to
        make technical modifications necessary to do so. The Licensor
        waives and/or agrees not to assert any right or authority to
        forbid You from making technical modifications necessary to
        exercise the Licensed Rights, including technical modifications
        necessary to circumvent Effective Technological Measures. For
        purposes of this Public License, simply making modifications
        authorized by this Section 2(a)(4) never produces Adapted
        Material.
    -   5. Downstream recipients.
        -   A. Offer from the Licensor – Licensed Material. Every
            recipient of the Licensed Material automatically receives an
            offer from the Licensor to exercise the Licensed Rights
            under the terms and conditions of this Public License.
        -   B. Additional offer from the Licensor – Adapted Material.
            Every recipient of Adapted Material from You automatically
            receives an offer from the Licensor to exercise the Licensed
            Rights in the Adapted Material under the conditions of the
            Adapter's License You apply.
        -   C. No downstream restrictions. You may not offer or impose
            any additional or different terms or conditions on, or apply
            any Effective Technological Measures to, the Licensed
            Material if doing so restricts exercise of the Licensed
            Rights by any recipient of the Licensed Material.
    -   6. No endorsement. Nothing in this Public License constitutes or
        may be construed as permission to assert or imply that You are,
        or that Your use of the Licensed Material is, connected with, or
        sponsored, endorsed, or granted official status by, the Licensor
        or others designated to receive attribution as provided in
        Section 3(a)(1)(A)(i).
-   b. Other rights.
    -   1. Moral rights, such as the right of integrity, are not
        licensed under this Public License, nor are publicity, privacy,
        and/or other similar personality rights; however, to the extent
        possible, the Licensor waives and/or agrees not to assert any
        such rights held by the Licensor to the limited extent necessary
        to allow You to exercise the Licensed Rights, but not otherwise.
    -   2. Patent and trademark rights are not licensed under this
        Public License.
    -   3. To the extent possible, the Licensor waives any right to
        collect royalties from You for the exercise of the Licensed
        Rights, whether directly or through a collecting society under
        any voluntary or waivable statutory or compulsory licensing
        scheme. In all other cases the Licensor expressly reserves any
        right to collect such royalties, including when the Licensed
        Material is used other than for NonCommercial purposes.

Section 3 – License Conditions.

Your exercise of the Licensed Rights is expressly made subject to the
following conditions.

-   a. Attribution.
    -   1. If You Share the Licensed Material (including in modified
        form), You must:
        -   A. retain the following if it is supplied by the Licensor
            with the Licensed Material:
            -   i. identification of the creator(s) of the Licensed
                Material and any others designated to receive
                attribution, in any reasonable manner requested by the
                Licensor (including by pseudonym if designated);
            -   ii. a copyright notice;
            -   iii. a notice that refers to this Public License;
            -   iv. a notice that refers to the disclaimer of
                warranties;
            -   v. a URI or hyperlink to the Licensed Material to the
                extent reasonably practicable;

        -   B. indicate if You modified the Licensed Material and retain
            an indication of any previous modifications; and
        -   C. indicate the Licensed Material is licensed under this
            Public License, and include the text of, or the URI or
            hyperlink to, this Public License.
    -   2. You may satisfy the conditions in Section 3(a)(1) in any
        reasonable manner based on the medium, means, and context in
        which You Share the Licensed Material. For example, it may be
        reasonable to satisfy the conditions by providing a URI or
        hyperlink to a resource that includes the required information.
    -   3. If requested by the Licensor, You must remove any of the
        information required by Section 3(a)(1)(A) to the extent
        reasonably practicable.
-   b. ShareAlike.In addition to the conditions in Section 3(a), if You
    Share Adapted Material You produce, the following conditions also
    apply.
    -   1. The Adapter's License You apply must be a Creative Commons
        license with the same License Elements, this version or later,
        or a BY-NC-SA Compatible License.
    -   2. You must include the text of, or the URI or hyperlink to, the
        Adapter's License You apply. You may satisfy this condition in
        any reasonable manner based on the medium, means, and context in
        which You Share Adapted Material.
    -   3. You may not offer or impose any additional or different terms
        or conditions on, or apply any Effective Technological Measures
        to, Adapted Material that restrict exercise of the rights
        granted under the Adapter's License You apply.

Section 4 – Sui Generis Database Rights.

Where the Licensed Rights include Sui Generis Database Rights that apply
to Your use of the Licensed Material:

-   a. for the avoidance of doubt, Section 2(a)(1) grants You the right
    to extract, reuse, reproduce, and Share all or a substantial portion
    of the contents of the database for NonCommercial purposes only;
-   b. if You include all or a substantial portion of the database
    contents in a database in which You have Sui Generis Database
    Rights, then the database in which You have Sui Generis Database
    Rights (but not its individual contents) is Adapted Material,
    including for purposes of Section 3(b); and
-   c. You must comply with the conditions in Section 3(a) if You Share
    all or a substantial portion of the contents of the database.
    For the avoidance of doubt, this Section 4 supplements and does not
    replace Your obligations under this Public License where the
    Licensed Rights include other Copyright and Similar Rights.

Section 5 – Disclaimer of Warranties and Limitation of Liability.

-   a. Unless otherwise separately undertaken by the Licensor, to the
    extent possible, the Licensor offers the Licensed Material as-is and
    as-available, and makes no representations or warranties of any kind
    concerning the Licensed Material, whether express, implied,
    statutory, or other. This includes, without limitation, warranties
    of title, merchantability, fitness for a particular purpose,
    non-infringement, absence of latent or other defects, accuracy, or
    the presence or absence of errors, whether or not known or
    discoverable. Where disclaimers of warranties are not allowed in
    full or in part, this disclaimer may not apply to You.
-   b. To the extent possible, in no event will the Licensor be liable
    to You on any legal theory (including, without limitation,
    negligence) or otherwise for any direct, special, indirect,
    incidental, consequential, punitive, exemplary, or other losses,
    costs, expenses, or damages arising out of this Public License or
    use of the Licensed Material, even if the Licensor has been advised
    of the possibility of such losses, costs, expenses, or damages.
    Where a limitation of liability is not allowed in full or in part,
    this limitation may not apply to You.
-   c. The disclaimer of warranties and limitation of liability provided
    above shall be interpreted in a manner that, to the extent possible,
    most closely approximates an absolute disclaimer and waiver of all
    liability.

Section 6 – Term and Termination.

-   a. This Public License applies for the term of the Copyright and
    Similar Rights licensed here. However, if You fail to comply with
    this Public License, then Your rights under this Public License
    terminate automatically.
-   b. Where Your right to use the Licensed Material has terminated
    under Section 6(a), it reinstates:

    -   1. automatically as of the date the violation is cured, provided
        it is cured within 30 days of Your discovery of the violation;
        or
    -   2. upon express reinstatement by the Licensor.

    For the avoidance of doubt, this Section 6(b) does not affect any
    right the Licensor may have to seek remedies for Your violations of
    this Public License.

-   c. For the avoidance of doubt, the Licensor may also offer the
    Licensed Material under separate terms or conditions or stop
    distributing the Licensed Material at any time; however, doing so
    will not terminate this Public License.
-   d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
    License.

Section 7 – Other Terms and Conditions.

-   a. The Licensor shall not be bound by any additional or different
    terms or conditions communicated by You unless expressly agreed.
-   b. Any arrangements, understandings, or agreements regarding the
    Licensed Material not stated herein are separate from and
    independent of the terms and conditions of this Public License.

Section 8 – Interpretation.

-   a. For the avoidance of doubt, this Public License does not, and
    shall not be interpreted to, reduce, limit, restrict, or impose
    conditions on any use of the Licensed Material that could lawfully
    be made without permission under this Public License.
-   b. To the extent possible, if any provision of this Public License
    is deemed unenforceable, it shall be automatically reformed to the
    minimum extent necessary to make it enforceable. If the provision
    cannot be reformed, it shall be severed from this Public License
    without affecting the enforceability of the remaining terms and
    conditions.
-   c. No term or condition of this Public License will be waived and no
    failure to comply consented to unless expressly agreed to by the
    Licensor.
-   d. Nothing in this Public License constitutes or may be interpreted
    as a limitation upon, or waiver of, any privileges and immunities
    that apply to the Licensor or You, including from the legal
    processes of any jurisdiction or authority.

Creative Commons is not a party to its public licenses. Notwithstanding,
Creative Commons may elect to apply one of its public licenses to
material it publishes and in those instances will be considered the
"Licensor." The text of the Creative Commons public licenses is
dedicated to the public domain under the CC0 Public Domain Dedication.
Except for the limited purpose of indicating that material is shared
under a Creative Commons public license or as otherwise permitted by the
Creative Commons policies published at creativecommons.org/policies,
Creative Commons does not authorize the use of the trademark "Creative
Commons" or any other trademark or logo of Creative Commons without its
prior written consent including, without limitation, in connection with
any unauthorized modifications to any of its public licenses or any
other arrangements, understandings, or agreements concerning use of
licensed material. For the avoidance of doubt, this paragraph does not
form part of the public licenses.

Creative Commons may be contacted at creativecommons.org.


================================================
FILE: LaMP/data/datasets.py
================================================
from torch.utils.data import Dataset
import json
import datasets
import torch

def get_all_labels(task):
    if task == "LaMP-1":
        return ["[1]","[2]"]
    elif task == "LaMP-2":
        return ['sci-fi', 'based on a book', 'comedy', 'action', 'twist ending', 'dystopia', 'dark comedy', 'classic', 'psychology', 'fantasy', 'romance', 'thought-provoking', 'social commentary', 'violence', 'true story']
    elif task == "LaMP-3":
        return ["1", "2", "3", "4", "5"]
    elif task == "LaMP-4":
        return []
    elif task == "LaMP-5":
        return []
    elif task == "LaMP-6":
        return []
    elif task == "LaMP-7":
        return []

def create_preprocessor(tokenizer, max_length):
    def preprocess_dataset(examples):
        inputs = [example for example in examples["source"]]
        targets = [example for example in examples["target"]]
        model_inputs = tokenizer(inputs, text_target=targets, max_length=max_length, truncation=True)
        return model_inputs
    return preprocess_dataset

def create_preprocessor_scores(tokenizer, max_length):
    def preprocess_dataset(examples):
        inputs = [example for example in examples["source"]]
        targets = [example for example in examples["target"]]
        model_inputs = tokenizer(inputs, text_target=targets, max_length=max_length, truncation=True)
        model_inputs['id_1'] = examples['id_1']
        model_inputs['id_2'] = examples['id_2']
        return model_inputs
    return preprocess_dataset

def create_preprocessor_scores_seq(tokenizer, max_length):
    def preprocess_dataset(examples):
        inputs = [example for example in examples["source"]]
        targets = [example for example in examples["target"]]
        model_inputs = tokenizer(inputs, text_target=targets, max_length=max_length, truncation=True)
        model_inputs['id'] = examples['id']
        return model_inputs
    return preprocess_dataset

def convert_to_hf_dataset(dataset, cache_dir):
    def gen():
        for idx in range(len(dataset)):
            yield dataset[idx]
    return datasets.Dataset.from_generator(gen, cache_dir = cache_dir)

class GeneralSeq2SeqDataset(Dataset):

    def __init__(self, data_addr, use_profile, task, create_prompt = None) -> None:
        super().__init__()
        with open(data_addr) as file:
            self.data = json.load(file)
        self.use_profile = use_profile
        self.task = task
        assert not (use_profile ^ (create_prompt != None)), "You should provide a prompt maker function when you use profile"
        self.create_prompt = create_prompt

    def __getitem__(self, index):
        if self.use_profile:
            return {
                "id" : self.data[index]['id'],
                "source" : self.create_prompt(self.data[index]['input'], self.data[index]['profile'], self.task),
                "target" : self.data[index]['output']
            }
        else:
            return {
                "id" : self.data[index]['id'],
                "source" : self.data[index]['input'],
                "target" : self.data[index]['output']
            }
    
    def __len__(self):
        return len(self.data)

class GeneralSeq2SeqForScoreGenerationDataset(Dataset):

    def __init__(self, data_addr, use_profile, task, create_prompt = None, max_prof_size = -1) -> None:
        super().__init__()
        with open(data_addr) as file:
            self.data = json.load(file)
        self.use_profile = use_profile
        self.task = task
        assert not (use_profile ^ (create_prompt != None)), "You should provide a prompt maker function when you use profile"
        self.create_prompt = create_prompt
        self.max_prof_size = max_prof_size
        self.size = 0
        self.index_dict = dict()
        for i, x in enumerate(self.data):
            for j, y in enumerate(x['profile']):
                if max_prof_size == -1 or j < self.max_prof_size:
                    self.index_dict[self.size] = (i, j)
                    self.size += 1

    def __getitem__(self, index):
        self.use_profile = True
        i, j = self.index_dict[index]
        if self.use_profile:
            return {
                "source" : self.create_prompt(self.data[i]['input'], [self.data[i]['profile'][j]], self.task),
                "target" : self.data[i]['output'],
                "id_1" : self.data[i]['id'],
                "id_2" : self.data[i]['profile'][j]['id']
            }
        else:
            return {
                "source" : self.data[index]['input'],
                "target" : self.data[index]['output']
            }
    
    def __len__(self):
        return self.size

================================================
FILE: LaMP/evaluate_llm.py
================================================
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, AutoModelForCausalLM
# from transformers.models.llama import LlamaTokenizer
from transformers.data.data_collator import DataCollatorForSeq2Seq
import argparse
from metrics.classification_metrics import create_metric_f1_accuracy, create_metric_mae_rmse
from metrics.generation_metrics import create_metric_bleu_rouge_meteor
from data.datasets import get_all_labels, GeneralSeq2SeqDataset, create_preprocessor, convert_to_hf_dataset
from prompts.prompts import create_prompt_generator
import json
import os

parser = argparse.ArgumentParser()

parser.add_argument("--validation_data", required = True)
parser.add_argument("--model_addr", required = True)
parser.add_argument("--task", required = True)
parser.add_argument("--output_dir", required = True)
parser.add_argument("--use_profile", action = "store_true")
parser.add_argument("--max_length", type = int, default = 256)
parser.add_argument("--generation_max_length", type = int, default = 128)
parser.add_argument("--per_device_batch_size", type = int, default = 16)
parser.add_argument("--generation_num_beams", type = int, default = 4)
parser.add_argument("--num_retrieved", type = int, default = 1)
parser.add_argument("--retriever", default = "bm25")
parser.add_argument("--is_ranked", action = "store_true")
parser.add_argument("--cache_dir", default = "./cache")


if __name__ == "__main__":

    opts = parser.parse_args()
    model = AutoModelForSeq2SeqLM.from_pretrained(opts.model_addr, cache_dir=opts.cache_dir)
    tokenizer = AutoTokenizer.from_pretrained(opts.model_addr, cache_dir=opts.cache_dir)
    collator = DataCollatorForSeq2Seq(tokenizer = tokenizer, model = model, max_length = opts.max_length)

    task = opts.task
    if opts.use_profile:
        prompt_generator, contriver = create_prompt_generator(opts.num_retrieved, opts.retriever, opts.is_ranked, opts.max_length, tokenizer)
    else:
        prompt_generator, contriver = None, None

    if task == "LaMP-1":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-2":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-3":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_mae_rmse(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-4":
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-5":
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-7":
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-6":
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
   
    eval_dataset = convert_to_hf_dataset(eval_dataset, cache_dir = opts.cache_dir).map(create_preprocessor(tokenizer = tokenizer, max_length = opts.max_length), batched=True)

    if contriver:
        contriver = contriver.to("cpu")

    training_args = Seq2SeqTrainingArguments(
        output_dir = opts.output_dir,
        do_eval = True,
        per_device_eval_batch_size = opts.per_device_batch_size,
        generation_num_beams = opts.generation_num_beams,
        predict_with_generate = True,
        eval_accumulation_steps = 1,
        generation_max_length = opts.generation_max_length
    )

    trainer = Seq2SeqTrainer(
        model = model,
        args = training_args,
        data_collator = collator,
        eval_dataset = eval_dataset,
        tokenizer = tokenizer,
        compute_metrics = compute_metrics
    )
    results = trainer.evaluate(eval_dataset)
    print(results)

    with open(os.path.join(opts.output_dir,'results_output.json'), 'w') as file:
        json.dump(results, file, indent = 4)

================================================
FILE: LaMP/metrics/classification_metrics.py
================================================
import numpy as np
import evaluate

def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [label.strip() for label in labels]

    return preds, labels

def create_metric_f1_accuracy(tokenizer, all_labels):
    f1_metric = evaluate.load("f1")
    accuracy_metric = evaluate.load("accuracy")
    def create_mapping(x):
        try:
            return all_labels.index(x)
        except:
            print(x)
            return -1
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
        decoded_preds = [create_mapping(x) for x in decoded_preds]
        decoded_labels = [create_mapping(x) for x in decoded_labels]
        result_acc = accuracy_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_f1 = f1_metric.compute(predictions=decoded_preds, references=decoded_labels, labels=list(range(len(all_labels))), average = "macro")
        result = {"accuracy" : result_acc["accuracy"], "f1" : result_f1["f1"]}
        return result
    return compute_metrics

def create_metric_f1_accuracy_bert(all_labels):
    f1_metric = evaluate.load("f1")
    accuracy_metric = evaluate.load("accuracy")
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        preds = np.argmax(preds, axis=1)
        result_acc = accuracy_metric.compute(predictions=preds, references=labels)
        result_f1 = f1_metric.compute(predictions=preds, references=labels, labels=list(range(len(all_labels))), average = "macro")
        result = {"accuracy" : result_acc["accuracy"], "f1" : result_f1["f1"]}
        return result
    return compute_metrics

def create_metric_mae_rmse_bert(all_labels):
    mse_metric = evaluate.load("mse")
    mae_metric = evaluate.load("mae")
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        preds = np.argmax(preds, axis=1)
        result_mae = mae_metric.compute(predictions=preds, references=labels)
        result_rmse = mse_metric.compute(predictions=preds, references=labels, squared = False)
        result = {"mae" : result_mae["mae"], "rmse" : result_rmse["mse"]}
        return result
    return compute_metrics

def create_metric_mae_rmse(tokenizer, all_labels):
    mse_metric = evaluate.load("mse")
    mae_metric = evaluate.load("mae")
    def create_mapping(x, y):
        try:
            return float(x)
        except:
            print(x)
            y = float(y)
            if abs(1 - y) > abs(5 - y):
                return 1.0
            else:
                return 5.0
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
        decoded_preds = [create_mapping(x,y) for x,y in zip(decoded_preds, decoded_labels)]
        decoded_labels = [create_mapping(x,x) for x in decoded_labels]
        result_mae = mae_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_rmse = mse_metric.compute(predictions=decoded_preds, references=decoded_labels, squared = False)
        result = {"mae" : result_mae["mae"], "rmse" : result_rmse["mse"]}
        return result
    return compute_metrics


def create_metric_f1_accuracy_chatgpt(all_labels):
    f1_metric = evaluate.load("f1")
    accuracy_metric = evaluate.load("accuracy")
    def create_mapping(x):
        try:
            return all_labels.index(x)
        except:
            print(x)
            return -1
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
        decoded_preds = [create_mapping(x) for x in decoded_preds]
        decoded_labels = [create_mapping(x) for x in decoded_labels]
        result_acc = accuracy_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_f1 = f1_metric.compute(predictions=decoded_preds, references=decoded_labels, labels=list(range(len(all_labels))), average = "macro")
        result = {"accuracy" : result_acc["accuracy"], "f1" : result_f1["f1"]}
        return result
    return compute_metrics

def create_metric_mae_rmse_chatgpt(all_labels):
    mse_metric = evaluate.load("mse")
    mae_metric = evaluate.load("mae")
    def create_mapping(x, y):
        try:
            return float(x)
        except:
            print(x)
            y = float(y)
            if abs(1 - y) > abs(5 - y):
                return 1.0
            else:
                return 5.0
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
        decoded_preds = [create_mapping(x,y) for x,y in zip(decoded_preds, decoded_labels)]
        decoded_labels = [create_mapping(x,x) for x in decoded_labels]
        result_mae = mae_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_rmse = mse_metric.compute(predictions=decoded_preds, references=decoded_labels, squared = False)
        result = {"mae" : result_mae["mae"], "rmse" : result_rmse["mse"]}
        return result
    return compute_metrics

================================================
FILE: LaMP/metrics/generation_metrics.py
================================================
import numpy as np
import evaluate

def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [[label.strip()] for label in labels]

    return preds, labels

def create_metric_bleu_rouge_meteor(tokenizer):
    bleu_metric = evaluate.load("sacrebleu")
    rouge_metric = evaluate.load('rouge')
    meteor_metric = evaluate.load('meteor')
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
        result_bleu = bleu_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_rouge = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_meteor = meteor_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result = {"bleu" : result_bleu["score"], "rouge-1" : result_rouge["rouge1"], "rouge-2" : result_rouge["rouge2"], "rouge-L" : result_rouge["rougeL"], "rouge-LSum" : result_rouge["rougeLsum"], "meteor" : result_meteor['meteor']}
        return result
    return compute_metrics

def create_metric_bleu_rouge_meteor_chatgpt():
    bleu_metric = evaluate.load("sacrebleu")
    rouge_metric = evaluate.load('rouge')
    meteor_metric = evaluate.load('meteor')
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
        result_bleu = bleu_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_rouge = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_meteor = meteor_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result = {"bleu" : result_bleu["score"], "rouge-1" : result_rouge["rouge1"], "rouge-2" : result_rouge["rouge2"], "rouge-L" : result_rouge["rougeL"], "rouge-LSum" : result_rouge["rougeLsum"], "meteor" : result_meteor['meteor']}
        return result
    return compute_metrics


================================================
FILE: LaMP/profile_item_utilization_scorer.py
================================================
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, AutoModelForCausalLM
# from transformers.models.llama import LlamaTokenizer
from transformers.data.data_collator import DataCollatorForSeq2Seq
import argparse
from metrics.classification_metrics import create_metric_f1_accuracy, create_metric_mae_rmse
from metrics.generation_metrics import create_metric_bleu_rouge_meteor
from data.datasets import get_all_labels, GeneralSeq2SeqForScoreGenerationDataset, create_preprocessor_scores, convert_to_hf_dataset
from prompts.prompts import create_prompt_generator
import tqdm
import datasets
import os
import json

parser = argparse.ArgumentParser()

parser.add_argument("--train_data", required = True)
parser.add_argument("--model_name", required = True)
parser.add_argument("--task", required = True)
parser.add_argument("--output_dir", required = True)
parser.add_argument("--max_length", type = int, default = 512)
parser.add_argument("--generation_max_length", type = int, default = 128)
parser.add_argument("--per_device_batch_size", type = int, default = 16)
parser.add_argument("--generation_num_beams", type = int, default = 4)
parser.add_argument("--cache_dir", default = "./cache")
parser.add_argument("--start_index", type = int, default=0)
parser.add_argument("--end_index", type = int, default=-1)
parser.add_argument("--profile_size", type = int,required = True)


if __name__ == "__main__":

    opts = parser.parse_args()
    model = AutoModelForSeq2SeqLM.from_pretrained(opts.model_name, cache_dir=opts.cache_dir)
    tokenizer = AutoTokenizer.from_pretrained(opts.model_name, cache_dir=opts.cache_dir)
    collator = DataCollatorForSeq2Seq(tokenizer = tokenizer, model = model, max_length = opts.max_length)

    task = opts.task
    prompt_generator, contriver = create_prompt_generator(1, "bm25", True, opts.max_length, tokenizer)
    
    if task == "LaMP-1":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqForScoreGenerationDataset(opts.train_data, True, task, prompt_generator, opts.profile_size)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-2":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqForScoreGenerationDataset(opts.train_data, True, task, prompt_generator, opts.profile_size)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-3":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqForScoreGenerationDataset(opts.train_data, True, task, prompt_generator, opts.profile_size)
        compute_metrics = create_metric_mae_rmse(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-4":
        eval_dataset = GeneralSeq2SeqForScoreGenerationDataset(opts.train_data, True, task, prompt_generator, opts.profile_size)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-5":
        eval_dataset = GeneralSeq2SeqForScoreGenerationDataset(opts.train_data, True, task, prompt_generator, opts.profile_size)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-7":
        eval_dataset = GeneralSeq2SeqForScoreGenerationDataset(opts.train_data, True, task, prompt_generator, opts.profile_size)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-6":
        eval_dataset = GeneralSeq2SeqForScoreGenerationDataset(opts.train_data, True, task, prompt_generator, opts.profile_size)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    
    eval_dataset = convert_to_hf_dataset(eval_dataset, opts.cache_dir).map(create_preprocessor_scores(tokenizer = tokenizer, max_length = opts.max_length), batched=True)

    if contriver:
        contriver = contriver.to("cpu")

    training_args = Seq2SeqTrainingArguments(
        output_dir = opts.output_dir,
        do_eval = True,
        per_device_eval_batch_size = 1,
        generation_num_beams = opts.generation_num_beams,
        predict_with_generate = True,
        eval_accumulation_steps = 1,
        generation_max_length = opts.generation_max_length
    )

    trainer = Seq2SeqTrainer(
        model = model,
        args = training_args,
        data_collator = collator,
        eval_dataset = eval_dataset,
        tokenizer = tokenizer,
        compute_metrics = compute_metrics
    )

    results_dict = dict()

    for i, x in enumerate(tqdm.tqdm(eval_dataset)):
        if i < opts.start_index:
            continue
        if i >= opts.end_index and opts.end_index != -1:
            break
        metrics = trainer.predict(datasets.Dataset.from_list([x])).metrics
        results_dict[f"{x['id_1']}-{x['id_2']}"] = {k.replace("test_", '') : v for k,v in metrics.items()}
    
    with open(os.path.join(opts.output_dir, f"scores_{opts.start_index}_{opts.end_index}.json"), "w") as file:
        json.dump(results_dict, file, indent = 4)

================================================
FILE: LaMP/prompts/contriever_retriever.py
================================================
import torch
from prompts.utils import batchify

def mean_pooling(token_embeddings, mask):
    token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.)
    sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]
    return sentence_embeddings

def retrieve_top_k_with_contriever(contriver, tokenizer, corpus, profile, query, k):
    query_tokens = tokenizer([query], padding=True, truncation=True, return_tensors='pt').to("cuda:0")
    output_query = contriver(**query_tokens)
    output_query = mean_pooling(output_query.last_hidden_state, query_tokens['attention_mask'])
    batch_size = 4
    scores = []
    batched_corpus = batchify(corpus, batch_size)
    for batch in batched_corpus:
        tokens_batch = tokenizer(batch, padding=True, truncation=True, return_tensors='pt').to("cuda:0")
        outputs_batch = contriver(**tokens_batch)
        outputs_batch = mean_pooling(outputs_batch.last_hidden_state, tokens_batch['attention_mask'])
        temp_scores = output_query.squeeze() @ outputs_batch.T
        scores.extend(temp_scores.tolist())
    topk_values, topk_indices = torch.topk(torch.tensor(scores), k)
    return [profile[m] for m in topk_indices.tolist()]


================================================
FILE: LaMP/prompts/prompts.py
================================================
from rank_bm25 import BM25Okapi
from transformers import AutoTokenizer, AutoModel
from prompts.utils import extract_strings_between_quotes, extract_after_article, extract_after_review, extract_after_paper, add_string_after_title, extract_after_colon, extract_after_description, extract_after_abstract
from prompts.contriever_retriever import retrieve_top_k_with_contriever
import random

def classification_citation_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["abstract"]}' for x in profile]
    extracted = extract_strings_between_quotes(inp)
    query = f'{extracted[1]} {extracted[2]}'
    return corpus, query

def classification_news_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["text"]}' for x in profile]
    query = extract_after_article(inp)
    return corpus, query

def classification_movies_query_corpus_maker(inp, profile):
    corpus = [f'{x["description"]}' for x in profile]
    query = extract_after_description(inp)
    return corpus, query

def classification_review_query_corpus_maker(inp, profile):
    corpus = [f'{x["text"]}' for x in profile]
    query = extract_after_review(inp)
    return corpus, query

def generation_news_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["text"]}' for x in profile]
    query = extract_after_article(inp)
    return corpus, query

def generation_paper_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["abstract"]}' for x in profile]
    query = extract_after_paper(inp)
    return corpus, query

def generation_paper_long_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["abstract"]}' for x in profile]
    query = extract_after_abstract(inp)
    return corpus, query


def parphrase_tweet_query_corpus_maker(inp, profile):
    corpus = [f'{x["text"]}' for x in profile]
    query = extract_after_colon(inp)
    return corpus, query

def generation_avocado_query_corpus_maker(inp, profile):
    corpus = [f'{x["text"]}' for x in profile]
    query = extract_after_colon(inp)
    return corpus, query

def generation_avocado_long_query_corpus_maker(inp, profile):
    corpus = [f'{x["text"]} {x["title"]}' for x in profile]
    query = extract_after_colon(inp)
    return corpus, query

def create_classification_citation_prompt(inp, profile, max_length, tokenizer):
    prompts = []
    per_p_max_length = (max_length - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    for p in profile:
        tokens = tokenizer(p["title"], max_length=per_p_max_length + saved_tokens - 2, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - 2
        new_title = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{new_title}"'
        prompts.append(prompt)
    return add_string_after_title(inp, ", and ".join(prompts))

def create_classification_news_prompt(inp, profile, max_length, tokenizer): # good
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'the category for the article: " " is "{p["category"]}" ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'the category for the article: "{new_text}" is "{p["category"]}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_classification_movies_prompt(inp, profile, max_length, tokenizer): # good
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'the tag for the movie: " " is "{p["tag"]}" ')['input_ids'])
        tokens = tokenizer(p["description"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'the tag for the movie: "{new_text}" is "{p["tag"]}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_classification_review_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'{p["score"]} is the score for " " ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'{p["score"]} is the score for "{new_text}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_generation_news_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is the title for " " ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is the title for "{new_text}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_generation_paper_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1) - len(tokenizer("Following the given patterns")['input_ids'])) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is a title " " ')['input_ids'])
        tokens = tokenizer(p["abstract"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_asbtract = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is a title for "{new_asbtract}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. Following the given patterns {inp}'

def create_generation_paper_long_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1) - len(tokenizer("Following the given patterns")['input_ids'])) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is the title " " ')['input_ids'])
        tokens = tokenizer(p["abstract"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_asbtract = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is the title for "{new_asbtract}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. Following the given patterns {inp}'


def create_parphrase_tweet_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1) - len(tokenizer("are written by user. Following the given patterns")['input_ids'])) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"" ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_asbtract = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{new_asbtract}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)} are written by a person. Following the given patterns {inp}'

def create_generation_avocado_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is the title for " " ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is the title for "{new_text}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_generation_avocado_long_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1) - len(tokenizer("are written by user. Following the given patterns")['input_ids'])) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is the title for " " ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is the title for "{new_text}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. Following the given patterns {inp}'

def create_prompt_generator(num_retrieve, ret_type = "bm25", is_ranked = False, max_length = 512, tokenizer = None):
    contriever = None
    if ret_type == "contriever" and not is_ranked:
        tokenizer = AutoTokenizer.from_pretrained('facebook/contriever')
        contriever = AutoModel.from_pretrained('facebook/contriever').to("cuda:0")
        contriever.eval()

    def prompt(inp, profile, task):
        if task == "LaMP-1":
            corpus, query = classification_citation_query_corpus_maker(inp, profile)
        elif task == "LaMP-2-old":
            corpus, query = classification_news_query_corpus_maker(inp, profile)
        elif task == "LaMP-2":
            corpus, query = classification_movies_query_corpus_maker(inp, profile)
        elif task == "LaMP-3":
            corpus, query = classification_review_query_corpus_maker(inp, profile)
        elif task == "LaMP-4":
            corpus, query = generation_news_query_corpus_maker(inp, profile)
        elif task == "LaMP-5":
            corpus, query = generation_paper_query_corpus_maker(inp, profile)
        elif task == "LaMP-7":
            corpus, query = parphrase_tweet_query_corpus_maker(inp, profile)
        elif task == "LaMP-6":
            corpus, query = generation_avocado_query_corpus_maker(inp, profile)
        
        if not is_ranked:
            if ret_type == "bm25":
                tokenized_corpus = [x.split() for x in corpus]
                bm25 = BM25Okapi(tokenized_corpus)
                tokenized_query = query.split()
                selected_profs = bm25.get_top_n(tokenized_query, profile, n=num_retrieve)
            elif ret_type == "contriever":
                selected_profs = retrieve_top_k_with_contriever(contriever, tokenizer, corpus, profile, query, num_retrieve)
            elif ret_type == "random":
                selected_profs = random.choices(profile, k = num_retrieve)
            elif ret_type == "recency":
                profile = sorted(profile, key=lambda x: tuple(map(int, str(x['date']).split("-"))))
                selected_profs = profile[-num_retrieve:][::-1]
        else:
            if ret_type == "recency_contriever":
                selected_profs_cont = profile[:num_retrieve // 2]
                profile = sorted(profile, key=lambda x: tuple(map(int, str(x['date']).split("-"))))
                selected_profs_rec = profile[-(num_retrieve // 2):][::-1]
                selected_profs = selected_profs_cont + selected_profs_rec
            else:
                selected_profs_cont = profile[:num_retrieve]
                selected_profs = selected_profs_cont
        factor = 0.6
        while True:
            try:
                max_len_prompt = max_length - min(len(tokenizer(inp)['input_ids']), int(factor * max_length))
                if task == "LaMP-1":
                    return create_classification_citation_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-2-old":
                    return create_classification_news_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-2":
                    return create_classification_movies_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-3":
                    return create_classification_review_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-4":
                    return create_generation_news_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-5":
                    return create_generation_paper_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-7":
                    return create_parphrase_tweet_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-6":
                    return create_generation_avocado_prompt(inp, selected_profs, max_len_prompt, tokenizer)
            except:
                factor -= 0.1
                if factor < 0:
                    print("not possible")
                    return inp
    return prompt, contriever

================================================
FILE: LaMP/prompts/utils.py
================================================
def extract_strings_between_quotes(input_string):
    output_list = []
    inside_quotes = False
    current_string = ''
    
    for char in input_string:
        if char == '"' and not inside_quotes:
            inside_quotes = True
        elif char == '"' and inside_quotes:
            inside_quotes = False
            output_list.append(current_string)
            current_string = ''
        elif inside_quotes:
            current_string += char
    
    return output_list

def extract_after_article(input_string):
    article_index = input_string.find('article:')
    if article_index == -1:
        return None
    return input_string[article_index + len('article:'):].strip()

def extract_after_description(input_string):
    article_index = input_string.find('description:')
    if article_index == -1:
        return None
    return input_string[article_index + len('description:'):].strip()


def extract_after_review(input_string):
    article_index = input_string.find('review:')
    if article_index == -1:
        return None
    return input_string[article_index + len('review:'):].strip()

def extract_after_paper(input_string):
    article_index = input_string.find('paper:')
    if article_index == -1:
        return None
    return input_string[article_index + len('paper:'):].strip()

def extract_after_abstract(input_string):
    article_index = input_string.find('abstract:')
    if article_index == -1:
        return None
    return input_string[article_index + len('abstract:'):].strip()

def extract_after_colon(input_string):
    article_index = input_string.find(':')
    if article_index == -1:
        return None
    return input_string[article_index + len(':'):].strip()


def add_string_after_title(original_string, string_to_add):
    title_index = original_string.find("title")
    
    if title_index == -1:
        return original_string
    
    return original_string[:title_index+5] + ", and " + string_to_add + original_string[title_index+5:]

def batchify(lst, batch_size):
    return [lst[i:i+batch_size] for i in range(0, len(lst), batch_size)]

================================================
FILE: LaMP/rank_profiles.py
================================================
import torch
from prompts.utils import batchify
from transformers import AutoModel, AutoTokenizer
import json
import tqdm
from prompts.utils import extract_strings_between_quotes, extract_after_article, extract_after_review, extract_after_paper, add_string_after_title, extract_after_colon, extract_after_abstract, extract_after_description
from rank_bm25 import BM25Okapi
import argparse

parser = argparse.ArgumentParser()

parser.add_argument("--input_data_addr", required = True)
parser.add_argument("--output_ranking_addr", required = True)
parser.add_argument("--task", required = True)
parser.add_argument("--ranker", required = True)
parser.add_argument("--batch_size", type = int, default=16)
parser.add_argument("--use_date", action='store_true')
parser.add_argument("--contriever_checkpoint", default="facebook/contriever")


def mean_pooling(token_embeddings, mask):
    token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.)
    sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]
    return sentence_embeddings

def retrieve_top_k_with_contriver(contriver, tokenizer, corpus, profile, query, k, batch_size = 16):
    query_tokens = tokenizer([query], padding=True, truncation=True, return_tensors='pt').to("cuda:0")
    output_query = contriver(**query_tokens)
    output_query = mean_pooling(output_query.last_hidden_state, query_tokens['attention_mask'])
    scores = []
    batched_corpus = batchify(corpus, batch_size)
    for batch in batched_corpus:
        tokens_batch = tokenizer(batch, padding=True, truncation=True, return_tensors='pt').to("cuda:0")
        outputs_batch = contriver(**tokens_batch)
        outputs_batch = mean_pooling(outputs_batch.last_hidden_state, tokens_batch['attention_mask'])
        temp_scores = output_query.squeeze() @ outputs_batch.T
        scores.extend(temp_scores.tolist())
    topk_values, topk_indices = torch.topk(torch.tensor(scores), k)
    return [profile[m] for m in topk_indices.tolist()]

def retrieve_top_k_with_bm25(corpus, profile, query, k):
    tokenized_corpus = [x.split() for x in corpus]
    bm25 = BM25Okapi(tokenized_corpus)
    tokenized_query = query.split()
    selected_profs = bm25.get_top_n(tokenized_query, profile, n=k)
    return selected_profs

def classification_citation_query_corpus_maker(inp, profile, use_date):
    if use_date:
        corpus = [f'{x["title"]} {x["abstract"]} date: {x["date"]}' for x in profile]
    else:
        corpus = [f'{x["title"]} {x["abstract"]}' for x in profile]
    ids = [x['id'] for x in profile]
    extracted = extract_strings_between_quotes(inp)
    query = f'{extracted[1]} {extracted[2]}'
    return corpus, query, ids

def classification_review_query_corpus_maker(inp, profile, use_date):
    if use_date:
        corpus = [f'{x["text"]} date: {x["date"]}' for x in profile]
    else:
        corpus = [f'{x["text"]}' for x in profile]
    ids = [x['id'] for x in profile]
    query = extract_after_review(inp)
    return corpus, query, ids

def generation_news_query_corpus_maker(inp, profile, use_date):
    if use_date:
        corpus = [f'{x["title"]} {x["text"]} date: {x["date"]}' for x in profile]
    else:
        corpus = [f'{x["title"]} {x["text"]}' for x in profile]
    ids = [x['id'] for x in profile]
    query = extract_after_article(inp)
    return corpus, query, ids

def generation_paper_query_corpus_maker(inp, profile, use_date):
    if use_date:
        corpus = [f'{x["title"]} {x["abstract"]} date: {x["date"]}' for x in profile]
    else:
        corpus = [f'{x["title"]} {x["abstract"]}' for x in profile]
    ids = [x['id'] for x in profile]
    query = extract_after_colon(inp)
    return corpus, query, ids

def parphrase_tweet_query_corpus_maker(inp, profile, use_date):
    if use_date:
        corpus = [f'{x["text"]} date: {x["date"]}' for x in profile]
    else:
        corpus = [f'{x["text"]}' for x in profile]
    query = extract_after_colon(inp)
    ids = [x['id'] for x in profile]
    return corpus, query, ids

def generation_avocado_query_corpus_maker(inp, profile, use_date):
    if use_date:
        corpus = [f'{x["text"]} date: {x["date"]}' for x in profile]
    else:
        corpus = [f'{x["text"]}' for x in profile]
    ids = [x['id'] for x in profile]
    query = extract_after_colon(inp)
    return corpus, query, ids

def classification_movies_query_corpus_maker(inp, profile, use_date):
    if use_date:
        corpus = [f'{x["description"]} date: {x["date"]}' for x in profile]
    else:
        corpus = [f'{x["description"]}' for x in profile]
    query = extract_after_description(inp)
    ids = [x['id'] for x in profile]
    return corpus, query, ids


if __name__ == "__main__":

    opts = parser.parse_args()
    task = opts.task
    ranker = opts.ranker

    with open(opts.input_data_addr) as file:
        dataset = json.load(file)
    
    rank_dict = dict()

    for data in tqdm.tqdm(dataset):
        inp = data['input']
        profile = data['profile']
        if task == "LaMP-1":
            corpus, query, ids = classification_citation_query_corpus_maker(inp, profile, opts.use_date)
        elif task == "LaMP-3":
            corpus, query, ids = classification_review_query_corpus_maker(inp, profile, opts.use_date)
        elif task == "LaMP-2":
            corpus, query = classification_movies_query_corpus_maker(inp, profile, opts.use_date)
        elif task == "LaMP-4":
            corpus, query, ids = generation_news_query_corpus_maker(inp, profile, opts.use_date)
        elif task == "LaMP-5":
            corpus, query, ids = generation_paper_query_corpus_maker(inp, profile, opts.use_date)
        elif task == "LaMP-7":
            corpus, query, ids = parphrase_tweet_query_corpus_maker(inp, profile, opts.use_date)
        elif task == "LaMP-6":
            corpus, query, ids = generation_avocado_query_corpus_maker(inp, profile, opts.use_date)
        
        if ranker == "contriever":
            tokenizer = AutoTokenizer.from_pretrained(opts.contriever_checkpoint)
            contriver = AutoModel.from_pretrained(opts.contriever_checkpoint).to("cuda:0")
            contriver.eval()
            randked_profile = retrieve_top_k_with_contriver(contriver, tokenizer, corpus, profile, query, len(profile), opts.batch_size)
        elif ranker == "bm25":
            randked_profile = retrieve_top_k_with_bm25(corpus, profile, query, len(profile))
        elif ranker == "recency":
            profile = sorted(profile, key=lambda x: tuple(map(int, str(x['date']).split("-"))))
            randked_profile = profile[::-1]

        data['profile'] = randked_profile

        rank_dict[data['id']] = [x['id'] for x in randked_profile]

    
    with open(opts.output_ranking_addr, "w") as file:
        json.dump(rank_dict, file)

================================================
FILE: LaMP/requirements.txt
================================================
mail_parser==3.15.0
numpy==1.24.2
rank_bm25==0.2.2
torch==2.0.0
tqdm==4.65.0
transformers==4.27.1


================================================
FILE: LaMP/retriever_utilization_scorer.py
================================================
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, AutoModelForCausalLM
# from transformers.models.llama import LlamaTokenizer
from transformers.data.data_collator import DataCollatorForSeq2Seq
import argparse
from metrics.classification_metrics import create_metric_f1_accuracy, create_metric_mae_rmse
from metrics.generation_metrics import create_metric_bleu_rouge_meteor
from data.datasets import get_all_labels, GeneralSeq2SeqDataset, create_preprocessor_scores_seq, convert_to_hf_dataset
from prompts.prompts import create_prompt_generator
import tqdm
import datasets
import os
import json

parser = argparse.ArgumentParser()

parser.add_argument("--data_addr", required = True)
parser.add_argument("--model_name", required = True)
parser.add_argument("--task", required = True)
parser.add_argument("--output_dir", required = True)
parser.add_argument("--use_profile", action = "store_true")
parser.add_argument("--max_length", type = int, default = 256)
parser.add_argument("--generation_max_length", type = int, default = 128)
parser.add_argument("--generation_num_beams", type = int, default = 4)
parser.add_argument("--num_retrieved", type = int, default = 4)
parser.add_argument("--retriever", default = "bm25")
parser.add_argument("--is_ranked", action = "store_true")
parser.add_argument("--cache_dir", default = "./cache")



if __name__ == "__main__":

    opts = parser.parse_args()
    model = AutoModelForSeq2SeqLM.from_pretrained(opts.model_name, cache_dir=opts.cache_dir)
    tokenizer = AutoTokenizer.from_pretrained(opts.model_name, cache_dir=opts.cache_dir)
    collator = DataCollatorForSeq2Seq(tokenizer = tokenizer, model = model, max_length = opts.max_length)

    task = opts.task
    if opts.use_profile:
        prompt_generator, contriver = create_prompt_generator(opts.num_retrieved, opts.retriever, opts.is_ranked, opts.max_length, tokenizer)
    else:
        prompt_generator, contriver = None, None

    if task == "LaMP-1":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.data_addr, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-2":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.data_addr, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-3":
        labels = get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.data_addr, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_mae_rmse(tokenizer = tokenizer, all_labels = labels)
    elif task == "LaMP-4":
        eval_dataset = GeneralSeq2SeqDataset(opts.data_addr, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-5":
        eval_dataset = GeneralSeq2SeqDataset(opts.data_addr, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-7":
        eval_dataset = GeneralSeq2SeqDataset(opts.data_addr, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
    elif task == "LaMP-6":
        eval_dataset = GeneralSeq2SeqDataset(opts.data_addr, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)

    eval_dataset = convert_to_hf_dataset(eval_dataset, cache_dir = opts.cache_dir).map(create_preprocessor_scores_seq(tokenizer = tokenizer, max_length = opts.max_length), batched=True)

    if contriver:
        contriver = contriver.to("cpu")

    training_args = Seq2SeqTrainingArguments(
        output_dir = opts.output_dir,
        do_eval = True,
        per_device_eval_batch_size = 1,
        generation_num_beams = opts.generation_num_beams,
        predict_with_generate = True,
        eval_accumulation_steps = 1,
        generation_max_length = opts.generation_max_length
    )

    trainer = Seq2SeqTrainer(
        model = model,
        args = training_args,
        data_collator = collator,
        eval_dataset = eval_dataset,
        tokenizer = tokenizer,
        compute_metrics = compute_metrics
    )

    results_dict = dict()
  
    for i, x in enumerate(tqdm.tqdm(eval_dataset)):
        preds = trainer.predict(datasets.Dataset.from_list([x]))
        metrics = preds.metrics
        output = tokenizer.batch_decode(preds.predictions, skip_special_tokens=True)[0].strip()
        results_dict[f"{x['id']}"] = {
            "metric" : {k.replace("test_", '') : v for k,v in metrics.items()}, 
            "output" : output, 
            "input":x['source']
        }

    with open(os.path.join(opts.output_dir, f"scores.json"), "w") as file:
        json.dump(results_dict, file, indent = 4)

================================================
FILE: LaMP/train_llm.py
================================================
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments
from transformers.data.data_collator import DataCollatorForSeq2Seq
import argparse
from metrics.classification_metrics import create_metric_f1_accuracy, create_metric_mae_rmse
from metrics.generation_metrics import create_metric_bleu_rouge_meteor
from data.datasets import get_all_labels, GeneralSeq2SeqDataset, create_preprocessor, convert_to_hf_dataset
from prompts.prompts import create_prompt_generator
import os
import json

parser = argparse.ArgumentParser()

parser.add_argument("--train_data", required = True)
parser.add_argument("--validation_data", required = True)
parser.add_argument("--test_data", default="")
parser.add_argument("--model_name", required = True)
parser.add_argument("--task", required = True)
parser.add_argument("--output_dir", required = True)
parser.add_argument("--retriever", default = "bm25")
parser.add_argument("--use_profile", action = "store_true")
parser.add_argument("--is_ranked", action = "store_true")
parser.add_argument("--max_length", type = int, default = 256)
parser.add_argument("--generation_max_length", type = int, default = 128)
parser.add_argument("--per_device_batch_size", type = int, default = 16)
parser.add_argument("--learning_rate", type = float, default = 5e-5)
parser.add_argument("--weight_decay", type = float, default = 0.0001)
parser.add_argument("--num_train_epochs", type = int, default = 10)
parser.add_argument("--lr_scheduler_type", default = "linear")
parser.add_argument("--warmup_ratio", type = float, default = 0.05)
parser.add_argument("--generation_num_beams", type = int, default = 4)
parser.add_argument("--num_retrieved", type = int, required=True)
parser.add_argument("--gradient_accumulation_steps", type = int, default = 1)
parser.add_argument("--cache_dir", default = "./cache")


if __name__ == "__main__":

    opts = parser.parse_args()
    
    model = AutoModelForSeq2SeqLM.from_pretrained(opts.model_name, cache_dir=opts.cache_dir)
    tokenizer = AutoTokenizer.from_pretrained(opts.model_name, cache_dir=opts.cache_dir)
    collator = DataCollatorForSeq2Seq(tokenizer = tokenizer, model = model, max_length = opts.max_length)

    task = opts.task
    if opts.use_profile:
        prompt_generator, contriver = create_prompt_generator(opts.num_retrieved, opts.retriever, opts.is_ranked, opts.max_length, tokenizer)
    else:
        prompt_generator, contriver = None, None

    greater_is_better = True
    if task == "LaMP-1":
        train_dataset, labels = GeneralSeq2SeqDataset(opts.train_data, opts.use_profile, task, prompt_generator), get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        if opts.test_data:
            test_dataset = GeneralSeq2SeqDataset(opts.test_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
        best_metric = "accuracy"
    elif task == "LaMP-2-old":
        train_dataset, labels = GeneralSeq2SeqDataset(opts.train_data, opts.use_profile, task, prompt_generator), get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        if opts.test_data:
            test_dataset = GeneralSeq2SeqDataset(opts.test_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
        best_metric = "accuracy"
    elif task == "LaMP-2":
        train_dataset, labels = GeneralSeq2SeqDataset(opts.train_data, opts.use_profile, task, prompt_generator), get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        if opts.test_data:
            test_dataset = GeneralSeq2SeqDataset(opts.test_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_f1_accuracy(tokenizer = tokenizer, all_labels = labels)
        best_metric = "accuracy"
    elif task == "LaMP-3":
        train_dataset, labels = GeneralSeq2SeqDataset(opts.train_data, opts.use_profile, task, prompt_generator), get_all_labels(task)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        if opts.test_data:
            test_dataset = GeneralSeq2SeqDataset(opts.test_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_mae_rmse(tokenizer = tokenizer, all_labels = labels)
        best_metric = "mae"
        greater_is_better = False
    elif task == "LaMP-4":
        train_dataset = GeneralSeq2SeqDataset(opts.train_data, opts.use_profile, task, prompt_generator)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        if opts.test_data:
            test_dataset = GeneralSeq2SeqDataset(opts.test_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
        best_metric = "rouge-1"
    elif task == "LaMP-5":
        train_dataset = GeneralSeq2SeqDataset(opts.train_data, opts.use_profile, task, prompt_generator)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        if opts.test_data:
            test_dataset = GeneralSeq2SeqDataset(opts.test_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
        best_metric = "rouge-1"
    elif task == "LaMP-7":
        train_dataset = GeneralSeq2SeqDataset(opts.train_data, opts.use_profile, task, prompt_generator)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        if opts.test_data:
            test_dataset = GeneralSeq2SeqDataset(opts.test_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
        best_metric = "rouge-1"
    elif task == "LaMP-6":
        train_dataset = GeneralSeq2SeqDataset(opts.train_data, opts.use_profile, task, prompt_generator)
        eval_dataset = GeneralSeq2SeqDataset(opts.validation_data, opts.use_profile, task, prompt_generator)
        if opts.test_data:
            test_dataset = GeneralSeq2SeqDataset(opts.test_data, opts.use_profile, task, prompt_generator)
        compute_metrics = create_metric_bleu_rouge_meteor(tokenizer = tokenizer)
        best_metric = "rouge-1"
    
    train_dataset = convert_to_hf_dataset(train_dataset, cache_dir = opts.cache_dir).map(create_preprocessor(tokenizer = tokenizer, max_length = opts.max_length), batched=True)
    eval_dataset = convert_to_hf_dataset(eval_dataset, cache_dir = opts.cache_dir).map(create_preprocessor(tokenizer = tokenizer, max_length = opts.max_length), batched=True)
    if opts.test_data:
        test_dataset = convert_to_hf_dataset(test_dataset, cache_dir = opts.cache_dir).map(create_preprocessor(tokenizer = tokenizer, max_length = opts.max_length), batched=True)

    if contriver:
        contriver = contriver.to("cpu")

    training_args = Seq2SeqTrainingArguments(
        output_dir = opts.output_dir,
        do_train = True,
        do_eval = True,
        evaluation_strategy = "epoch",
        per_device_train_batch_size = opts.per_device_batch_size,
        per_device_eval_batch_size = opts.per_device_batch_size,
        gradient_accumulation_steps = opts.gradient_accumulation_steps,
        learning_rate = opts.learning_rate,
        weight_decay = opts.weight_decay,
        num_train_epochs = opts.num_train_epochs,
        lr_scheduler_type = opts.lr_scheduler_type,
        warmup_ratio = opts.warmup_ratio,
        generation_num_beams = opts.generation_num_beams,
        predict_with_generate = True,
        save_strategy = "epoch",
        logging_steps = 50,
        eval_accumulation_steps = 1,
        generation_max_length = opts.generation_max_length,
        load_best_model_at_end = True,
        metric_for_best_model = best_metric,
        greater_is_better = greater_is_better
    )

    trainer = Seq2SeqTrainer(
        model = model,
        args = training_args,
        data_collator = collator,
        train_dataset = train_dataset,
        eval_dataset = eval_dataset,
        tokenizer = tokenizer,
        compute_metrics = compute_metrics
    )

    trainer.train()

    if opts.test_data:
        results = trainer.evaluate(test_dataset)
        print(results)

        with open(os.join(opts.output_dir,'results_output.json'), 'w') as file:
            json.dump(results, file, indent = 4)

================================================
FILE: LaMP/utils/merge_with_rank.py
================================================
import json 
import argparse

def merge(inps, outs, ranks):
    for inp in inps:
        for o in outs:
            if o['id'] == inp['id']:
                output = o['output']
                break
        new_profile = []
        for x in ranks[inp['id']]:
            for y in inp['profile']:
                if y['id'] == x:
                    new_profile.append(y)
                    break
        inp['profile'] = new_profile
        inp['output'] = output
    return inps

parser = argparse.ArgumentParser()

parser.add_argument("--lamp_questions_addr", required = True)
parser.add_argument("--lamp_output_addr", required = True)
parser.add_argument("--merged_output_addr", required = True)
parser.add_argument("--profile_ranking_addr", default="")

if __name__ == "__main__":
    opts = parser.parse_args()
    q_addr = opts.lamp_questions_addr
    o_addr = opts.lamp_output_addr
    rank_addr = opts.profile_ranking_addr
    res_addr = opts.merged_output_addr

    with open(q_addr) as qfile:
        inp = json.load(qfile)
    with open(o_addr) as oflie:
        out = json.load(oflie)
    if rank_addr:
        with open(rank_addr) as rflie:
            rank = json.load(rflie)
    else:
        rank = dict()
        for data in inp:
            rank[data['id']] = []
            for item in data['profile']:
                rank[data['id']].append(item['id'])

    with open(res_addr, "w") as resfile:
        res = merge(inp, out, rank)
        json.dump(res, resfile, indent=4)




================================================
FILE: PEFT/data/datasets.py
================================================
from torch.utils.data import Dataset
import json
import datasets
import torch
import random
from itertools import combinations

def sublists_between_2_and_k(lst, k):
    sublists = []
    for size in range(2, k+1):  # Iterate through sizes from 2 to k
        for comb in combinations(lst, size):
            sublists.append(list(comb))
    return sublists

def sample_sublists(lst, k, num_samples):
    sublists = []
    for i in range(k+1, len(lst)):
        sub = list(random.sample(lst[:i], k-1))
        sub.sort(key=lambda x: x['date'])
        sub += [lst[i]]
        sublists.append(sub)
    while len(sublists) < num_samples:
        idx = random.randint(k+1, len(lst) - 1)
        sub = list(random.sample(lst[:idx], k-1))
        sub.sort(key=lambda x: x['date'])
        sub += [lst[idx]]
        sublists.append(sub)
    return sublists

def get_all_labels(task):
    if task == "classification_citation":
        return ["[1]","[2]"]
    elif task == "classification_news":
        return ["food & drink", "sports", "education", "parents", "religion", "travel", "business", "crime", "science & technology", "culture & arts", "entertainment", "politics", "women", "style & beauty", "healthy living"]
    elif task == "classification_movies":
        return ['sci-fi', 'based on a book', 'comedy', 'action', 'twist ending', 'dystopia', 'dark comedy', 'classic', 'psychology', 'fantasy', 'romance', 'thought-provoking', 'social commentary', 'violence', 'true story']
    elif task == "classification_review":
        return ["1", "2", "3", "4", "5"]
    elif task == "generation_news":
        return []
    elif task == "generation_paper":
        return []
    elif task == "paraphrase_paper":
        return []

def create_preprocessor(tokenizer, max_length):
    def preprocess_dataset(examples):
        inputs = [example for example in examples["source"]]
        targets = [example for example in examples["target"]]
        model_inputs = tokenizer(inputs, text_target=targets, max_length=max_length, truncation=True)
        return model_inputs
    return preprocess_dataset

def convert_to_hf_dataset(dataset, cache_dir):
    def gen():
        for idx in range(len(dataset)):
            yield dataset[idx]
    return datasets.Dataset.from_generator(gen, cache_dir = cache_dir)

def create_input_output_gen_func(task):
    if task == "LaMP-1":
        def func(item):
            inp = f"Write an abstract for this title: {item['title']}"
            out = f'{item["abstract"]}'
            return inp, out
    elif task == "LaMP-2":
        def func(item):
            inp = f"Which tag does this movie relate to among the following tags? Just answer with the tag name without further explanation. tags: [sci-fi, based on a book, comedy, action, twist ending, dystopia, dark comedy, classic, psychology, fantasy, romance, thought-provoking, social commentary, violence, true story] description: {item['description']}"
            out = f'{item["tag"]}'
            return inp, out
    elif task == "LaMP-3":
        def func(item):
            inp = f"What is the score of the following review on a scale of 1 to 5? just answer with 1, 2, 3, 4, or 5 without further explanation. review: {item['text']}"
            out = f'{item["score"]}'
            return inp, out
    elif task == "LaMP-4":
        def func(item):
            inp = f"Generate a headline for the following article: {item['text']}"
            out = f'{item["title"]}'
            return inp, out
    elif task == "LaMP-5":
        def func(item):
            inp = f"Generate a title for the following abstract of a paper: {item['abstract']}"
            out = f'{item["title"]}'
            return inp, out
    elif task == "LaMP-6":
        def func(item):
            inp = f"Generate a subject for the following email: {item['text']}"
            out = f'{item["title"]}'
            return inp, out
    elif task == "LaMP-7":
        def func(item):
            percent = random.uniform(0.1, 0.25)
            tweet_words = item['text'].split()
            index = int(len(tweet_words) * percent)
            in_inp = " ".join(tweet_words[:index])
            in_out = " ".join(tweet_words[index:])
            inp = f"Complete the following tweet: {in_inp}"
            out = f'{in_out}'
            return inp, out
    return func


def create_per_user_dataset(data_addr, user_ids, task, cache_dir):
    with open(data_addr) as file:
        orig_dataset = json.load(file)
    seen_users = set()
    datasets = dict()
    input_output_gen_func = create_input_output_gen_func(task)
    for data in orig_dataset:
        uid = str(data['user_id'])
        if user_ids is not None and uid not in user_ids:
            continue
        if uid in seen_users:
            continue
        else:
            seen_users.add(uid)
        cur_dataset = []
        for i, item in enumerate(data['profile']):
            id = f'{uid}-{data["id"]}-{i}'
            inp, out = input_output_gen_func(item)
            cur_dataset.append(
                {
                    "id" : id,
                    "input" : inp,
                    "output" : out
                }
            )
        datasets[uid] = convert_to_hf_dataset(GeneralSeq2SeqDataset(cur_dataset), cache_dir)
    return datasets

def create_per_user_dataset_test(data_addr, user_ids, task, cache_dir):
    with open(data_addr) as file:
        orig_dataset = json.load(file)
    seen_users = set()
    datasets = dict()
    for data in orig_dataset:
        uid = str(data['user_id'])
        if user_ids is not None and uid not in user_ids:
            continue
        elif uid not in seen_users:
            seen_users.add(uid)
            datasets[uid] = []
        
        datasets[uid].append(
            {
                "id" : data["id"],
                "input" : data["input"],
                "output" : data["output"]
            }
        )
    
    for key, value in datasets.items():
        datasets[key] = convert_to_hf_dataset(GeneralSeq2SeqDataset(value), cache_dir)
    return datasets

class GeneralSeq2SeqDataset(Dataset):

    def __init__(self, data) -> None:
        super().__init__()
        self.data = data

    def __getitem__(self, index):
        return {
            "id" : self.data[index]['id'],
            "source" : self.data[index]['input'],
            "target" : self.data[index]['output']
        }
    
    def __len__(self):
        return len(self.data)

================================================
FILE: PEFT/evaluate_llm.py
================================================
import argparse
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments
from transformers.data.data_collator import DataCollatorForSeq2Seq
from data.datasets import create_per_user_dataset_test, create_preprocessor
import json
from peft import get_peft_config, get_peft_model, LoraConfig, TaskType
import os
import torch
import glob
import numpy as np


parser = argparse.ArgumentParser()

parser.add_argument("--test_data", required = True)
parser.add_argument("--user_ids", default = "")
parser.add_argument("--golds_addr", required = True)
parser.add_argument("--task", required = True)
parser.add_argument("--user_checkpoints", required = True)
parser.add_argument("--output_dir", required = True)
parser.add_argument("--max_length", type = int, default = 512)
parser.add_argument("--num_shards", type = int, default = 1)
parser.add_argument("--shard_id", type = int, default = 0)
parser.add_argument("--generation_max_length", type = int, default = 128)
parser.add_argument("--per_device_batch_size", type = int, default = 16)
parser.add_argument("--generation_num_beams", type = int, default = 4)
parser.add_argument("--cache", default="./cache")

if __name__ == "__main__":
    opts = parser.parse_args()
    print(opts)
    with open(opts.user_ids) as file:
        all_user_ids = [str(x) for x in json.load(file)]
        shard_size = len(all_user_ids) // opts.num_shards + 1
        user_ids = all_user_ids[int(opts.shard_id * shard_size):int((opts.shard_id + 1) * shard_size)]
    
    user_datasets = create_per_user_dataset_test(opts.test_data, user_ids, opts.task, opts.cache)

    model_name_or_path = "google/flan-t5-xxl"
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path, cache_dir = opts.cache)
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, cache_dir = opts.cache)  
    processor = create_preprocessor(tokenizer = tokenizer, max_length = opts.max_length)
    collator = DataCollatorForSeq2Seq(tokenizer = tokenizer, model = model, max_length = opts.max_length)
    peft_config = LoraConfig(
        task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1, target_modules=['q','v','k']
    )
    model = get_peft_model(model, peft_config)

    final_outputs = []

    for key, dataset in user_datasets.items():
        model.unload()
        checkpoitns = glob.glob(os.path.join(opts.user_checkpoints, 'adaptors', key, '*'))
        if len(checkpoitns) > 0:
            checkpoint_addr = checkpoitns[0]
            print(checkpoint_addr)
            model.load_adapter(checkpoint_addr, key)
            model.set_adapter(key)
        encoded_dataset = dataset.map(processor, batched=True)
        
        training_args = Seq2SeqTrainingArguments(
            output_dir = opts.output_dir,
            do_train = False,
            do_eval = True,
            per_device_train_batch_size = opts.per_device_batch_size,
            generation_max_length = opts.generation_max_length,
            generation_num_beams = opts.generation_num_beams,
            predict_with_generate=True,
            eval_accumulation_steps = 1
        )

        trainer = Seq2SeqTrainer(
            model = model,
            args = training_args,
            data_collator = collator,
            train_dataset = encoded_dataset,
            tokenizer = tokenizer
        )

        preds = trainer.predict(encoded_dataset).predictions
        preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
        preds = tokenizer.batch_decode(preds, skip_special_tokens = True)
        for data, pred in zip(dataset, preds):
            final_outputs.append(
                {
                    "id" : data['id'],
                    "output" : pred
                }
            )
    prediction_addr = os.path.join(opts.output_dir, 'predictions.json')

    with open(prediction_addr, 'w') as file:
        json.dump(
            {
                "task" : opts.task,
                "golds" : final_outputs
            },
            file,
            indent=4
        )

================================================
FILE: PEFT/requirements.txt
================================================
datasets==2.8.0
regex==2022.10.31
sentencepiece==0.1.97
tokenizers==0.11.1
torch==2.0.1
tqdm==4.64.1
transformers==4.28.0
evaluate
absl-py
rouge-score
peft

================================================
FILE: PEFT/train_peft.py
================================================
import argparse
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments
from transformers.data.data_collator import DataCollatorForSeq2Seq
from data.datasets import create_per_user_dataset, create_preprocessor
import json
from peft import get_peft_config, get_peft_model, LoraConfig, TaskType
import os
import torch


def is_directory_empty(path):
    if os.path.exists(path):
        if not os.listdir(path):
            return True
        else:
            return False
    else:
        return True


parser = argparse.ArgumentParser()

parser.add_argument("--train_data", required = True)
parser.add_argument("--user_ids", default="")
parser.add_argument("--task", required = True)
parser.add_argument("--output_dir", required = True)
parser.add_argument("--max_length", type = int, default = 512)
parser.add_argument("--num_shards", type = int, default = 1)
parser.add_argument("--shard_id", type = int, default = 0)
parser.add_argument("--generation_max_length", type = int, default = 512)
parser.add_argument("--per_device_batch_size", type = int, default = 16)
parser.add_argument("--learning_rate", type = float, default = 5e-5)
parser.add_argument("--weight_decay", type = float, default = 0.0001)
parser.add_argument("--num_train_epochs", type = int, default = 30)
parser.add_argument("--lora_r", type = int, default = 8)
parser.add_argument("--lr_scheduler_type", default = "linear")
parser.add_argument("--warmup_ratio", type = float, default = 0.05)
parser.add_argument("--gradient_accumulation_steps", type = int, default = 1)
parser.add_argument("--cache", default="./cache")



if __name__ == "__main__":

    opts = parser.parse_args()
    print(opts)
    if opts.user_ids:
        with open(opts.user_ids) as file:
            all_user_ids = [str(x) for x in json.load(file)]
            shard_size = len(all_user_ids) // opts.num_shards + 1
            user_ids = all_user_ids[int(opts.shard_id * shard_size):int((opts.shard_id + 1) * shard_size)]
    else:
        user_ids = None
    user_datasets = create_per_user_dataset(opts.train_data, user_ids, opts.task, opts.cache)

    model_name_or_path = "google/flan-t5-xxl"
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path, cache_dir = opts.cache)
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, cache_dir = opts.cache)  
    processor = create_preprocessor(tokenizer = tokenizer, max_length = opts.max_length)
    collator = DataCollatorForSeq2Seq(tokenizer = tokenizer, model = model, max_length = opts.max_length)
    for key, dataset in user_datasets.items():
        print(key)
        if not is_directory_empty(os.path.join(opts.output_dir, 'adaptors', key)):
            continue
        peft_config = LoraConfig(
            task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=opts.lora_r, lora_alpha=32, lora_dropout=0.1, target_modules=['q','v','k']
        )
        model = get_peft_model(model, peft_config)
        encoded_dataset = dataset.map(processor, batched=True)
        training_args = Seq2SeqTrainingArguments(
            output_dir = os.path.join(opts.output_dir, 'adaptors', key),
            do_train = True,
            do_eval = False,
            per_device_train_batch_size = opts.per_device_batch_size,
            gradient_accumulation_steps = opts.gradient_accumulation_steps,
            learning_rate = opts.learning_rate,
            weight_decay = opts.weight_decay,
            num_train_epochs = opts.num_train_epochs,
            lr_scheduler_type = opts.lr_scheduler_type,
            warmup_ratio = opts.warmup_ratio,
            save_strategy = "epoch",
            save_total_limit=1,
            logging_steps = 10,
            generation_max_length = opts.generation_max_length,
            save_only_model = True
        )

        trainer = Seq2SeqTrainer(
            model = model,
            args = training_args,
            data_collator = collator,
            train_dataset = encoded_dataset,
            tokenizer = tokenizer
        )

        trainer.train()
        model.unload()
        trainer = None
        torch.cuda.empty_cache()
        
        

================================================
FILE: README.md
================================================
# Codes for papers on Large Language Models Personalization (LaMP)

[LaMP: When Large Language Models Meet Personalization](https://arxiv.org/abs/2304.11406)

This paper highlights the importance of personalization in the current state of natural language understanding and generation and introduces the LaMP benchmark --- a novel benchmark for training and evaluating language models for producing personalized outputs. LaMP offers a comprehensive evaluation framework with diverse language tasks and multiple entries for each user profile. It consists of seven personalized tasks, spanning across three classification and four text generation tasks. We further propose a retrieval augmentation approach that retrieves personalized items from user profiles to construct personalized prompts for large language models. The experiments conducted to establish fine-tuned and zero-shot baseline results for the benchmark conclude that LMs utilizing profile augmentation outperform their counterparts that do not factor in profile information.

```
@misc{salemi2023lamp,
      title={La{MP}: When Large Language Models Meet Personalization}, 
      author={Alireza Salemi and Sheshera Mysore and Michael Bendersky and Hamed Zamani},
      year={2023},
      eprint={2304.11406},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
```

[Optimization Methods for Personalizing Large Language Models through Retrieval Augmentation](https://arxiv.org/abs/2404.05970)

This paper studies retrieval-augmented approaches for personalizing large language models (LLMs), which potentially have a substantial impact on various applications and domains. We propose the first attempt to optimize the retrieval models that deliver a limited number of personal documents to large language models for the purpose of personalized generation. We develop two optimization algorithms that solicit feedback from the downstream personalized generation tasks for retrieval optimization--one based on reinforcement learning whose reward function is defined using any arbitrary metric for personalized generation and another based on knowledge distillation from the downstream LLM to the retrieval model. This paper also introduces a pre- and post-generation retriever selection model that decides what retriever to choose for each LLM input. Extensive experiments on diverse tasks from the language model personalization (LaMP) benchmark reveal statistically significant improvements in six out of seven datasets.

```
@misc{salemi2024optimization,
      title={Optimization Methods for Personalizing Large Language Models through Retrieval Augmentation}, 
      author={Alireza Salemi and Surya Kallumadi and Hamed Zamani},
      year={2024},
      eprint={2404.05970},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
```

[Comparing Retrieval-Augmentation and Parameter-Efficient Fine-Tuning for Privacy-Preserving Personalization of Large Language Models](https://arxiv.org/abs/2409.09510)

Privacy-preserving methods for personalizing large language models (LLMs) are relatively under-explored. There are two schools of thought on this topic: (1) generating personalized outputs by personalizing the input prompt through retrieval augmentation from the user's personal information (RAG-based methods), and (2) parameter-efficient fine-tuning of LLMs per user that considers efficiency and space limitations (PEFT-based methods). This paper presents the first systematic comparison between two approaches on a wide range of personalization tasks using seven diverse datasets. Our results indicate that RAG-based and PEFT-based personalization methods on average yield 14.92% and 1.07% improvements over the non-personalized LLM, respectively. We find that combining RAG with PEFT elevates these improvements to 15.98%. Additionally, we identify a positive correlation between the amount of user data and PEFT's effectiveness, indicating that RAG is a better choice for cold-start users (i.e., user's with limited personal data).

```
@misc{salemi2024comparingretrievalaugmentationparameterefficientfinetuning,
      title={Comparing Retrieval-Augmentation and Parameter-Efficient Fine-Tuning for Privacy-Preserving Personalization of Large Language Models}, 
      author={Alireza Salemi and Hamed Zamani},
      year={2024},
      eprint={2409.09510},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2409.09510}, 
}
```

## Data

You can download all the datasets from the links provided [here](https://lamp-benchmark.github.io/download). However, we provided the minimal ids to generate the dataset using our codes for the Personalized Email Subject Generation because this dataset is not publicly accessible. Follow the following section to generate that dataset.

### LaMP 6: Personalized Email Subject Generation (Avocado dataset)

The [Avocado](https://catalog.ldc.upenn.edu/LDC2015T03) dataset is not publicly accessible. However, we provided the samples' id and the code we used to generate our dataset. Therefore, if you get access to the dataset, you can quickly generate the dataset with the same format as the other datasets in LaMP using the following code:

```
python data/avocado/create_avocado_dataset.py \
    --avocado_files_dir \*Address to the directory containing zip files for avocado dataset 'avocado-1.0.2/data/text'*\ \
    --extract_addr \*A temp dir to extract the files for creating dataset*\ \
    --output_dir \*The directory to generate the final dataset*\ \
    --input_question_file_train \*The address to the train_questions.json file we provided in LaMP*\ \
    --input_question_file_dev \*The address to the dev_questions.json file we provided in LaMP*\ \
    --input_question_file_test \*The address to the test_questions.json file we provided in LaMP*\
```

## Evaluation

The instructions for evaluating your results on the test set are provided [here](https://lamp-benchmark.github.io/leaderboard). In order to evaluate your results on the dev set, we provided an evaluation script that can be found here:


Evaluate all tasks together:

```
python eval/eval_all.py \
    --golds_zip /*Address to all gold labels for all tasks zipped in a file*/ \
    --preds_zip /*Address to all predictions for all tasks zipped in a file*/ \
    --temp_dir /*Address to a temp dir for extracting files*/ \
    --output_file /*Address to the results file*/ \
```

Evaluate one task:

```
python eval/eval_task.py \
    --golds_json /*Address to gold labels for the task as a json file*/ \
    --preds_json /*Address to predictions for the task as a json file*/ \
    --task_name /*Name of the task [LaMP_1, LaMP_2, LaMP_3, LaMP_4, LaMP_5, LaMP_6, LaMP_7]*/
    --output_file /*Address to the results file*/ \
```

The pred files should follow the exact same format as the gold files:

```
{
    "task" : "/*task name*/",
    "golds" : [
        {
            "id" : "/*sample 1 id*/",
            "output" : "/*output of the model for the first sample*/"
        },
        ...,
        {
            "id" : "/*sample n id*/",
            "output" : "/*output of the model for the n'th sample*/"
        }
    ]
}
```

## Personalizing LLMs with RAG (LaMP)

You first need to create an environment for this using the following script:

```
python3 -m venv lamp_venv
source lamp_venv/bin/activate
pip install -r LaMP/requirements.txt
```

### Ranking Profiles based on the Input

The first step is to sort items in each user profile based on the input for the task:

```
cd LaMP
python rank_profiles.py \
    --input_data_addr /*input questions for one of the LaMP tasks*/ \
    --output_ranking_addr /*output address for the generated ranking file*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
    --ranker /*the ranking model to be used [bm25, contriever, recency]*/ \
    [optional] --use_date /*the batch size for ranking*/ \
    [optional] --use_date \ /*if used, it adds time to the text of each profile item*/
    [optional] --contriever_checkpoint /*address to the Contriever checkpoint to be used*/ \

```

After that, use the following script to sort the profiles in the dataset based on the ranking file:

```
cd LaMP
python utils/merge_with_rank.py \
    --lamp_questions_addr /*address to the LaMP task inputs file*/ \
    --lamp_output_addr /*address to the LaMP task outputs file*/ \
    --profile_ranking_addr /*address to the generated ranking file from the previous script*/
    --merged_output_addr /*address to the sorted dataset using the provided ranking file*/ \

```

### Training LLM with RAG

The next step is to train the LLM on a LaMP task:

```
cd LaMP
python train_llm.py \
    --train_data /*address to sorted training data using the previous step*/ \
    --validation_data /*address to sorted validation data using the previous step*/ \
    [optional] --test_data /*address to sorted test data using the previous step*/ \
    --model_name /*address to the model that should be used for initialization of the LLM*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
    --output_dir /*output directory to save results and checkpoints*/ \
    --retriever /*the ranking model to be used [bm25, contriever, recency]*/ \
    --use_profile \ /*used to perfrom personalization with RAG */
    --is_ranked \ /*used if you pre-ranked the profiles based on the provided retrieval model*/
    --num_retrieved /*number of items to be retrieved from the user profile*/ \ 
```

### Zero-shot Evaluation of LLM with RAG

You can also evaluate the LLMs with the following script:

```
cd LaMP
python evaluate_llm.py \
    --validation_data /*address to sorted validation data using the previous step*/ \
    --model_addr /*address to the model that should be used for initialization of the LLM*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
    --output_dir /*output directory to save results */ \
    --use_profile \ /*used to perfrom personalization with RAG */
    --retriever /*the ranking model to be used [bm25, contriever, recency]*/ \
    --is_ranked \ /*used if you pre-ranked the profiles based on the provided retrieval model*/
    --num_retrieved /*number of items to be retrieved from the user profile*/ \ 
```

## Optimizing Retrieval Model for Personalizing LLMs (ROPG)

This code uses the feedback from LLM to train a retrieval model for personalizing the LLM. You first need to create an environment for this using the following script:

```
python3 -m venv ropg_venv
source ropg_venv/bin/activate
pip install -r ROPG/requirements.txt
```

### Feedback Generation using LLM for Items in the User Profile

The first step is to collect feedback from the LLM using the following script:

```
cd LaMP
python profile_item_utilization_scorer.py \
    --train_data /*address to sorted training data using the previous steps*/ \
    --model_name /*address to the model that should be used for feedback generation*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
    --output_dir /*output directory to save results */ \
    --profile_size /*number of top k items from user profile to get feedback for them*/
```

### Optimizing Retrieval Model

You can use the following code to train a retrieval model based on the feedback generated from the previous step.

For training with ROPG-KD, which uses knowledge distillation, use the following script:

```
cd ROPG
NGPU=/*Number of GPUs*/ python -m torch.distributed.launch --nproc_per_node=/*Number of GPUs*/ train_kd.py \
    --train_data /*address to sorted training data using the previous steps*/ \
    --do_train \
    --scores_path /*address to the feedback file generated in the previous step*/
    --name /*output directory*/ \
    --ctx_size /*number of documents to be used for training the retrieval model for each query*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
    --temperature /*temperature for distillation*/
```

For training with ROPG-RL, which uses reinforcement learning, use the following script:

```
cd ROPG
NGPU=/*Number of GPUs*/ python -m torch.distributed.launch --nproc_per_node=/*Number of GPUs*/ train_rl.py \
    --train_data /*address to sorted training data using the previous steps*/ \
    --do_train \
    --scores_path /*address to the feedback file generated in the previous step*/
    --name /*output directory*/ \
    --ctx_size /*number of documents to be used for training the retrieval model for each query*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
```

## Retrieval Model Selection for Personalizing LLMs (RSPG)

This section uses the feedback from the LLM based on the performance of different retrieval models to train a retrieval model selector. You first need to create an environment for this using the following script:

```
python3 -m venv rspg_venv
source rspg_venv/bin/activate
pip install -r RSPG/requirements.txt
```


### Feedback Generation using LLM for each Retrieval Model

use the following code to get the feedback for each retrieval model in the retrieval model pool:

```
cd LaMP
python retriever_utilization_scorer.py \
    --data_addr /**address to sorted task data using the previous steps**/
    --model_name /*address to the model that should be used for feedback generation*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
    --output_dir /*output directory to save results */ \
    --use_profile \ /*use only in the case you want the feedback from RAG approach, shouldn't be used when getting feedback from an LLM without RAG*/
    --num_retrieved /*number of items to be retrieved from the user profile*/ \ 
    --retriever /*the retriever model that should be used to get feedback for*/ \
    --is_ranked \ /*used if you pre-ranked the profiles based on the provided retrieval model*/
```

You should use the following script with all the retrieval models in your retrieval model pool. In our paper we used Contriever, ROPG-RL, ROPG-KD, Recency, BM25, and no retrieval (no RAG).

### Optimizing Retrieval Model Selector

The first step is to combine all the feebacks got from the previous step and make a training and validation set:

```
cd RSPG
python utils/create_data.py \
    --retrivers_data_addr "/*address to feedback 1*/" "/*address to feedback 2*/" ... "/*address to feedback n*/" \
    --task_inputs_addr /*input questions for one of the LaMP tasks*/ \
    --task_outputs_addr /*outputs for one of the LaMP tasks*/ \
    --output_dataset_addr /*address to save the created dataset*/
    --metric /*the metric name that should be used as feedback [accuracy, rouge-1, rouge-l]*/ \
```

After this, you can use the following script to train the retrieval selector model (RSPG):

```
cd RSPG
NGPU=/*Number of GPUs*/ python -m torch.distributed.launch --nproc_per_node=/*Number of GPUs*/ rspg.py \
    --train_data /*address to the training data created in the previous step*/ \
    --val_data /*address to the validation data created in the previous step*/ \
    --rspg_type /*retrieval selection mode: [Pre, Post]*/
    --val_lamp_golds /*address to the a LaMP task output file for validation set*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
    --do_train \
    --name /*output directory*/ \
    --temperature /*distillation temperature*/ \
```

### Inference with Retrieval Model Selector

In order to do inference with RSPG, you can use the following script:

```
cd RSPG
NGPU=/*Number of GPUs*/ python -m torch.distributed.launch --nproc_per_node=/*Number of GPUs*/ rspg.py \
    --train_data /*address to the training data created in the previous step*/ \
    --val_data /*address to the validation data created in the previous step*/ \
    --rspg_type /*retrieval selection mode: [Pre, Post]*/
    --val_lamp_golds /*address to the a LaMP task output file for validation set*/ \
    --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
    --do_validation \
    --name /*output directory*/ \
    --model_path /*address to the checkpoint to be evaluated*/ \
```

## PEFT for Personalizing LLMs

This section trains an LLM per user on its personal data usign LoRA. You first need to create an environment for this using the following script:

```
python3 -m venv peft_venv
source peft_venv/bin/activate
pip install -r PEFT/requirements.txt
```

### Training LLM using LoRA on personal data
In order to train a LoRA adaptor per user, you can use the following script:

```
cd PEFT
python train_peft.py \
      --train_data /*address to sorted training data using the previous step*/ \
      --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
      --output_dir /*output directory to save per user checkpoints*/ \
      --lora_r /*lora r parameter*/
```

### inference using per user LLM
In order to inference using the LoRA adaptor per user, you can use the following script:

```
cd PEFT
python evaluate_llm.py \
      --test_data /*address to sorted test/validation data using the previous step*/ \
      --task /*name of the task [LaMP-1, LaMP-2, ..., LaMP-7]*/ \
      --output_dir /*output directory to save outputs*/ \
      --user_checkpoints /*directory containing per user checkpoints*/
```

## Reference

If you find this repository helpful, please cite the following works!

[LaMP: When Large Language Models Meet Personalization](https://arxiv.org/abs/2304.11406)

```
@misc{salemi2023lamp,
      title={La{MP}: When Large Language Models Meet Personalization}, 
      author={Alireza Salemi and Sheshera Mysore and Michael Bendersky and Hamed Zamani},
      year={2023},
      eprint={2304.11406},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
```

[Optimization Methods for Personalizing Large Language Models through Retrieval Augmentation](https://arxiv.org/abs/2404.05970)

```
@misc{salemi2024optimization,
      title={Optimization Methods for Personalizing Large Language Models through Retrieval Augmentation}, 
      author={Alireza Salemi and Surya Kallumadi and Hamed Zamani},
      year={2024},
      eprint={2404.05970},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
```

[Comparing Retrieval-Augmentation and Parameter-Efficient Fine-Tuning for Privacy-Preserving Personalization of Large Language Models](https://arxiv.org/abs/2409.09510)

```
@misc{salemi2024comparingretrievalaugmentationparameterefficientfinetuning,
      title={Comparing Retrieval-Augmentation and Parameter-Efficient Fine-Tuning for Privacy-Preserving Personalization of Large Language Models}, 
      author={Alireza Salemi and Hamed Zamani},
      year={2024},
      eprint={2409.09510},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2409.09510}, 
}
```

## License
LaMP (codes and data creation methods) is licensed by Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0). See the [CC-BY-NC-SA-4.0.txt](CC-BY-NC-SA-4.0.txt) file for details. For the datasets in this benchmark, you should follow their license.

## Acknowledgments

This work was supported in part by the Center for Intelligent Information Retrieval, in part by NSF grant #2143434, in part by the Office of Naval Research contract number N000142212688, and in part by Lowe's, in part by an Amazon Research Award, Fall 2022 CFP, in part by an award from Google, and in part by an award from Microsoft. Any opinions, findings and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect those of the sponsors.


================================================
FILE: ROPG/data/collators.py
================================================
from typing import Any, List
import numpy as np
import torch
import json


class ReaderToRetreieverCollator:

    def __init__(self, tokenizer, query_max_lenght, document_max_length, number_of_ctx, scores_addr = "") -> None:
        self.tokenizer = tokenizer
        self.query_max_lenght = query_max_lenght
        self.document_max_length = document_max_length
        self.number_of_ctx = number_of_ctx
        self.scores_addr = scores_addr
        if scores_addr:
            with open(scores_addr) as file:
                self.scores = json.load(file)
    
    def __call__(self, examples: List[dict]):
        query_tokens = self.tokenizer([x['query'] for x in examples], max_length = self.query_max_lenght, padding = True, return_tensors = 'pt', truncation = True)
        docs = []
        for x in examples:
            temp_docs = []
            for y in x['documents'][:self.number_of_ctx]:
                temp_docs.append(y)
            while len(temp_docs) < self.number_of_ctx:
                temp_docs.append("")
            docs.append(temp_docs)
        
        profiles = []
        for x in examples:
            temp_docs = []
            for y in x['profile'][:self.number_of_ctx]:
                temp_docs.append(y)
            profiles.append(temp_docs)
        
        documents_tokens = self.tokenizer([y for x in docs for y in x], max_length = self.document_max_length, padding = True, return_tensors = 'pt', truncation = True)
        documents_tokens['input_ids'] = documents_tokens['input_ids'].view(len(examples), self.number_of_ctx, -1)
        documents_tokens['token_type_ids'] = documents_tokens['token_type_ids'].view(len(examples), self.number_of_ctx, -1)
        documents_tokens['attention_mask'] = documents_tokens['attention_mask'].view(len(examples), self.number_of_ctx, -1)
        
        scores_batch = []
        for x in examples:
            scores_sample = []
            for prof in x['profile'][:self.number_of_ctx]:
                score = self.scores[f'{x["qid"]}-{prof["id"]}']
                scores_sample.append(score)
            scores_batch.append(scores_sample)
        # print(scores_batch)
        target_txt = [x['target'] for x in examples]       
        ctxs = torch.tensor([[len(x['documents'][:self.number_of_ctx])] for x in examples])
        return {
            "query_input_ids" : query_tokens['input_ids'],
            "query_token_type_ids" : query_tokens['token_type_ids'],
            "query_attention_mask" : query_tokens['attention_mask'],
            "documents_input_ids" : documents_tokens['input_ids'],
            "documents_token_type_ids" : documents_tokens['token_type_ids'],
            "documents_attention_mask" : documents_tokens['attention_mask'],
            "documents_ctxs" : ctxs,
            "batch_docs_text" : profiles,
            "batch_questions_text" : [x['query_raw'] for x in examples],
            "target_txt" : target_txt,
            "scores_gold" : scores_batch
        }


================================================
FILE: ROPG/data/datasets.py
================================================
from torch.utils.data import Dataset
import json
import datasets

def get_all_labels(task):
    if task == "LaMP-1":
        return ["[1]","[2]"]
    elif task == "LaMP-2":
        return ['sci-fi', 'based on a book', 'comedy', 'action', 'twist ending', 'dystopia', 'dark comedy', 'classic', 'psychology', 'fantasy', 'romance', 'thought-provoking', 'social commentary', 'violence', 'true story']
    elif task == "LaMP-3":
        return ["1", "2", "3", "4", "5"]
    else:
        return []

def create_preprocessor(tokenizer, max_length):
    def preprocess_dataset(examples):
        inputs = [example for example in examples["source"]]
        targets = [example for example in examples["target"]]
        model_inputs = tokenizer(inputs, text_target=targets, max_length=max_length, truncation=True)
        return model_inputs
    return preprocess_dataset

def create_preprocessor_chatgpt(tokenizer, max_length):
    def preprocess_dataset(examples):
        inputs = [example for example in examples["source"]]
        targets = [example for example in examples["target"]]
        model_inputs = tokenizer(inputs, text_target=targets, max_length=max_length, truncation=True)
        model_inputs = tokenizer.batch_decode(model_inputs['input_ids'], skip_special_tokens=True)
        return {"chatgpt_inputs" : model_inputs}
    return preprocess_dataset

def convert_to_hf_dataset(dataset):
    def gen():
        for idx in range(len(dataset)):
            yield dataset[idx]
    return datasets.Dataset.from_generator(gen)

class GeneralSeq2SeqDataset(Dataset):

    def __init__(self, data_addr, use_profile, task, create_prompt = None) -> None:
        super().__init__()
        with open(data_addr) as file:
            self.data = json.load(file)
        self.use_profile = use_profile
        self.task = task
        assert not (use_profile ^ (create_prompt != None)), "You should provide a prompt maker function when you use profile"
        self.create_prompt = create_prompt

    def __getitem__(self, index):
        if self.use_profile:
            return {
                "source" : self.create_prompt(self.data[index]['input'], self.data[index]['profile'], self.task),
                "target" : self.data[index]['output']
            }
        else:
            return {
                "source" : self.data[index]['input'],
                "target" : self.data[index]['output']
            }
    
    def __len__(self):
        return len(self.data)

class ReaderToRetrieverDataset(Dataset):

    def __init__(self, data_addr, task, create_query_corpus, is_llama = False) -> None:
        super().__init__()
        with open(data_addr) as file:
            self.data = json.load(file)
        self.task = task
        self.create_query_corpus = create_query_corpus
        self.is_llama = is_llama

    def __getitem__(self, index):
        query, corpus = self.create_query_corpus(self.data[index]['input'], self.data[index]['profile'])
        return {
            "qid" : self.data[index]['id'],
            "query_raw" : self.data[index]['input'] + " answer:" if self.is_llama else "",
            "query" : query,
            "documents" : corpus,
            "profile" : self.data[index]['profile'],
            "target" : self.data[index]['output']
        }
    
    def __len__(self):
        return len(self.data)

================================================
FILE: ROPG/models/optim.py
================================================
import torch

class WarmupLinearScheduler(torch.optim.lr_scheduler.LambdaLR):
    def __init__(self, optimizer, warmup_steps, scheduler_steps, min_ratio, fixed_lr, last_epoch=-1):
        self.warmup_steps = warmup_steps
        self.scheduler_steps = scheduler_steps
        self.min_ratio = min_ratio
        self.fixed_lr = fixed_lr
        super(WarmupLinearScheduler, self).__init__(
            optimizer, self.lr_lambda, last_epoch=last_epoch
        )

    def lr_lambda(self, step):
        if step < self.warmup_steps:
            return (1 - self.min_ratio)*step/float(max(1, self.warmup_steps)) + self.min_ratio

        if self.fixed_lr:
            return 1.0

        return max(0.0,
            1.0 + (self.min_ratio - 1) * (step - self.warmup_steps)/float(max(1.0, self.scheduler_steps - self.warmup_steps)),
        )


class FixedScheduler(torch.optim.lr_scheduler.LambdaLR):
    
    def __init__(self, optimizer, last_epoch=-1):
        super(FixedScheduler, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)
   
    def lr_lambda(self, step):
        return 1.0


def set_optim(opt, model):
    if opt.optim == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr)
    elif opt.optim == 'adamw':
        optimizer = torch.optim.AdamW(model.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)
    if opt.scheduler == 'fixed':
        scheduler = FixedScheduler(optimizer)
    elif opt.scheduler == 'linear':
        if opt.scheduler_steps is None:
            scheduler_steps = opt.total_steps
        else:
            scheduler_steps = opt.scheduler_steps
        scheduler = WarmupLinearScheduler(optimizer, warmup_steps=opt.warmup_steps, scheduler_steps=scheduler_steps, min_ratio=0., fixed_lr=opt.fixed_lr)
    return optimizer, scheduler

================================================
FILE: ROPG/models/retriever.py
================================================
from typing import Any
from transformers import BertModel
from transformers.configuration_utils import PretrainedConfig
import torch

class Contriever(BertModel):
    def __init__(self, config, pooling="average", **kwargs):
        super().__init__(config, add_pooling_layer=False)
        if not hasattr(config, "pooling"):
            self.config.pooling = pooling

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        normalize=False,
    ):

        model_output = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        last_hidden = model_output["last_hidden_state"]
        last_hidden = last_hidden.masked_fill(~attention_mask[..., None].bool(), 0.0)

        if self.config.pooling == "average":
            emb = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
        elif self.config.pooling == "cls":
            emb = last_hidden[:, 0]

        if normalize:
            emb = torch.nn.functional.normalize(emb, dim=-1)
        return emb





================================================
FILE: ROPG/prompts/contriever_retriever.py
================================================
import torch
from prompts.utils import batchify

def mean_pooling(token_embeddings, mask):
    token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.)
    sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]
    return sentence_embeddings

def retrieve_top_k_with_contriever(contriver, tokenizer, corpus, profile, query, k):
    query_tokens = tokenizer([query], padding=True, truncation=True, return_tensors='pt').to("cuda:0")
    output_query = contriver(**query_tokens)
    output_query = mean_pooling(output_query.last_hidden_state, query_tokens['attention_mask'])
    batch_size = 4
    scores = []
    batched_corpus = batchify(corpus, batch_size)
    for batch in batched_corpus:
        tokens_batch = tokenizer(batch, padding=True, truncation=True, return_tensors='pt').to("cuda:0")
        outputs_batch = contriver(**tokens_batch)
        outputs_batch = mean_pooling(outputs_batch.last_hidden_state, tokens_batch['attention_mask'])
        temp_scores = output_query.squeeze() @ outputs_batch.T
        scores.extend(temp_scores.tolist())
    topk_values, topk_indices = torch.topk(torch.tensor(scores), k)
    return [profile[m] for m in topk_indices.tolist()]


================================================
FILE: ROPG/prompts/prompts.py
================================================
from rank_bm25 import BM25Okapi
from transformers import AutoTokenizer, AutoModel
from prompts.utils import extract_strings_between_quotes, extract_after_article, extract_after_review, extract_after_paper, add_string_after_title, extract_after_colon, extract_after_abstract, extract_after_description
from prompts.contriever_retriever import retrieve_top_k_with_contriever
import random

def classification_citation_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["abstract"]} date: {x["date"]}' for x in profile]
    extracted = extract_strings_between_quotes(inp)
    query = f'{extracted[1]} {extracted[2]}'
    return corpus, query

def classification_news_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["text"]} date: {x["date"]}' for x in profile]
    query = extract_after_article(inp)
    return corpus, query

def classification_movies_query_corpus_maker(inp, profile):
    corpus = [f'{x["description"]} date: {x["date"]}' for x in profile]
    query = extract_after_description(inp)
    return corpus, query

def classification_review_query_corpus_maker(inp, profile):
    corpus = [f'{x["text"]} date: {x["date"]}' for x in profile]
    query = extract_after_review(inp)
    return corpus, query

def generation_news_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["text"]} date: {x["date"]}' for x in profile]
    query = extract_after_article(inp)
    return corpus, query

def generation_paper_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["abstract"]} date: {x["date"]}' for x in profile]
    query = extract_after_paper(inp)
    return corpus, query

def generation_paper_long_query_corpus_maker(inp, profile):
    corpus = [f'{x["title"]} {x["abstract"]} date: {x["date"]}' for x in profile]
    query = extract_after_abstract(inp)
    return corpus, query


def parphrase_tweet_query_corpus_maker(inp, profile):
    corpus = [f'{x["text"]} date: {x["date"]}' for x in profile]
    query = extract_after_colon(inp)
    return corpus, query

def generation_avocado_query_corpus_maker(inp, profile):
    corpus = [f'{x["text"]} date: {x["date"]}' for x in profile]
    query = extract_after_colon(inp)
    return corpus, query

def generation_avocado_long_query_corpus_maker(inp, profile):
    corpus = [f'{x["text"]} {x["title"]} date: {x["date"]}' for x in profile]
    query = extract_after_colon(inp)
    return corpus, query

def create_classification_citation_prompt(inp, profile, max_length, tokenizer):
    prompts = []
    per_p_max_length = (max_length - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    for p in profile:
        tokens = tokenizer(p["title"], max_length=per_p_max_length + saved_tokens - 2, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - 2
        new_title = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{new_title}"'
        prompts.append(prompt)
    return add_string_after_title(inp, ", and ".join(prompts))

def create_classification_news_prompt(inp, profile, max_length, tokenizer): # good
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'the category for the article: " " is "{p["category"]}" ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'the category for the article: "{new_text}" is "{p["category"]}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_classification_review_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'{p["score"]} is the score for " " ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'{p["score"]} is the score for "{new_text}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_generation_news_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is the title for " " ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is the title for "{new_text}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_generation_paper_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1) - len(tokenizer("Following the given patterns")['input_ids'])) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is a title " " ')['input_ids'])
        tokens = tokenizer(p["abstract"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_asbtract = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is a title for "{new_asbtract}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. Following the given patterns {inp}'

def create_generation_paper_long_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1) - len(tokenizer("Following the given patterns")['input_ids'])) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is the title " " ')['input_ids'])
        tokens = tokenizer(p["abstract"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_asbtract = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is the title for "{new_asbtract}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. Following the given patterns {inp}'

def create_parphrase_tweet_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1) - len(tokenizer("are written by user. Following the given patterns")['input_ids'])) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"" ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_asbtract = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{new_asbtract}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)} are written by a person. Following the given patterns {inp}'

def create_generation_avocado_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is the title for " " ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is the title for "{new_text}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_generation_avocado_long_prompt(inp, profile, max_length, tokenizer):
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1) - len(tokenizer("are written by user. Following the given patterns")['input_ids'])) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'"{p["title"]}" is the title for " " ')['input_ids'])
        tokens = tokenizer(p["text"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'"{p["title"]}" is the title for "{new_text}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. Following the given patterns {inp}'

def create_classification_movies_prompt(inp, profile, max_length, tokenizer): # good
    per_p_max_length = (max_length - 1 - 2 * (len(profile) - 1)) // len(profile)
    saved_tokens = 0
    prompts = []
    for p in profile:
        needed_part_len = len(tokenizer(f'the tag for the movie: " " is "{p["tag"]}" ')['input_ids'])
        tokens = tokenizer(p["description"], max_length=per_p_max_length + saved_tokens - needed_part_len, truncation=True)
        saved_tokens += per_p_max_length - len(tokens['input_ids']) - needed_part_len
        new_text = tokenizer.batch_decode([tokens['input_ids']], skip_special_tokens=True)[0]
        prompt = f'the tag for the movie: "{new_text}" is "{p["tag"]}" '
        prompts.append(prompt)
    return f'{", and ".join(prompts)}. {inp}'

def create_query_corpus_generator(task):
    def create_query_corpus(inp, profile):
        if task == "LaMP-1":
            corpus, query = classification_citation_query_corpus_maker(inp, profile)
        elif task == "LaMP-2":
            corpus, query = classification_movies_query_corpus_maker(inp, profile)
        elif task == "LaMP-3":
            corpus, query = classification_review_query_corpus_maker(inp, profile)
        elif task == "LaMP-4":
            corpus, query = generation_news_query_corpus_maker(inp, profile)
        elif task == "LaMP-5":
            corpus, query = generation_paper_query_corpus_maker(inp, profile)
        elif task == "LaMP-7":
            corpus, query = parphrase_tweet_query_corpus_maker(inp, profile)
        elif task == "LaMP-6":
            corpus, query = generation_avocado_query_corpus_maker(inp, profile)
        return query, corpus
    return create_query_corpus

def create_prompt_generator(num_retrieve, ret_type = "bm25", is_ranked = False, max_length = 512, tokenizer = None):
    contriever = None
    if ret_type == "contriever" and not is_ranked:
        tokenizer = AutoTokenizer.from_pretrained('facebook/contriever')
        contriever = AutoModel.from_pretrained('facebook/contriever').to("cuda:0")
        contriever.eval()

    def prompt(inp, profile, task):
        if task == "LaMP-1":
            corpus, query = classification_citation_query_corpus_maker(inp, profile)
        elif task == "LaMP-2":
            corpus, query = classification_movies_query_corpus_maker(inp, profile)
        elif task == "LaMP-3":
            corpus, query = classification_review_query_corpus_maker(inp, profile)
        elif task == "LaMP-4":
            corpus, query = generation_news_query_corpus_maker(inp, profile)
        elif task == "LaMP-5":
            corpus, query = generation_paper_query_corpus_maker(inp, profile)
        elif task == "LaMP-7":
            corpus, query = parphrase_tweet_query_corpus_maker(inp, profile)
        elif task == "LaMP-6":
            corpus, query = generation_avocado_query_corpus_maker(inp, profile)
        
        if not is_ranked:
            if ret_type == "bm25":
                tokenized_corpus = [x.split() for x in corpus]
                bm25 = BM25Okapi(tokenized_corpus)
                tokenized_query = query.split()
                selected_profs = bm25.get_top_n(tokenized_query, profile, n=num_retrieve)
            elif ret_type == "contriever":
                selected_profs = retrieve_top_k_with_contriever(contriever, tokenizer, corpus, profile, query, num_retrieve)
            elif ret_type == "random":
                selected_profs = random.choices(profile, k = num_retrieve)
            elif ret_type == "rec":
                selected_profs = profile[-num_retrieve:][::-1]
        else:
            selected_profs = profile[:num_retrieve]
        factor = 0.6
        while True:
            try:
                max_len_prompt = max_length - min(len(tokenizer(inp)['input_ids']), int(factor * max_length))
                if task == "LaMP-1":
                    return create_classification_citation_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-3":
                    return create_classification_review_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-2":
                    return create_classification_movies_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-4":
                    return create_generation_news_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-5":
                    return create_generation_paper_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-7":
                    return create_parphrase_tweet_prompt(inp, selected_profs, max_len_prompt, tokenizer)
                elif task == "LaMP-6":
                    return create_generation_avocado_prompt(inp, selected_profs, max_len_prompt, tokenizer)
            except:
                factor -= 0.1
                if factor < 0:
                    print(len(profile))
                    print(len(selected_profs))
                    return inp
                    # raise RuntimeError("not possible")
    return prompt, contriever

================================================
FILE: ROPG/prompts/utils.py
================================================
import re

def extract_strings_between_quotes(input_string):
    pattern = r'"(.*?)"'
    titles = re.findall(pattern, input_string)
    return titles

def extract_after_article(input_string):
    article_index = input_string.find('article:')
    if article_index == -1:
        return None
    return input_string[article_index + len('article:'):].strip()

def extract_after_description(input_string):
    article_index = input_string.find('description:')
    if article_index == -1:
        return None
    return input_string[article_index + len('description:'):].strip()

def extract_after_review(input_string):
    article_index = input_string.find('review:')
    if article_index == -1:
        return None
    return input_string[article_index + len('review:'):].strip()

def extract_after_paper(input_string):
    article_index = input_string.find('paper:')
    if article_index == -1:
        return None
    return input_string[article_index + len('paper:'):].strip()

def extract_after_colon(input_string):
    article_index = input_string.find(':')
    if article_index == -1:
        return None
    return input_string[article_index + len(':'):].strip()

def extract_after_abstract(input_string):
    article_index = input_string.find('abstract:')
    if article_index == -1:
        return None
    return input_string[article_index + len('abstract:'):].strip()


def add_string_after_title(original_string, string_to_add):
    title_index = original_string.find("title")
    
    if title_index == -1:
        return original_string
    
    return original_string[:title_index+5] + ", and " + string_to_add + original_string[title_index+5:]

def batchify(lst, batch_size):
    return [lst[i:i+batch_size] for i in range(0, len(lst), batch_size)]

================================================
FILE: ROPG/requirements.txt
================================================
evaluate==0.4.0
numpy==1.24.3
rank_bm25==0.2.2
torch==2.0.1
transformers==4.29.2
rouge_score

================================================
FILE: ROPG/train_kd.py
================================================
from pathlib import Path
from utils.distributed import init_distributed_mode, init_signal_handler
import torch
import numpy as np
import os
from torch.utils.data import DataLoader, DistributedSampler, SequentialSampler, RandomSampler
import tqdm
import argparse
from transformers import AutoModel, AutoTokenizer, AutoModelForSeq2SeqLM
from data.collators import ReaderToRetreieverCollator
from data.datasets import ReaderToRetrieverDataset, get_all_labels
from prompts.prompts import create_query_corpus_generator
from trainers.trainer import KDReaderToRetrieverTrainer
from models.retriever import Contriever
from utils.util import average_main
from utils.log import init_logger
from models.optim import set_optim
from utils.util import save_checkpoint, load_checkpoint


def train(opts, model, optimizer, scheduler, step, dataset, collator, checkpoint_path):
    
    if opts.is_main:
        try:
            tb_logger = torch.utils.tensorboard.SummaryWriter(Path(opts.checkpoint_dir)/opts.name)
        except:
           tb_logger = None
           logger.warning('Tensorboard is not available.')

    torch.manual_seed(opts.global_rank + opts.seed)
    if opts.is_distributed:
        train_sampler = DistributedSampler(dataset, num_replicas=opts.n_gpu_per_node, rank=opts.local_rank)
    else:
        train_sampler = RandomSampler(dataset)
    bar = tqdm.tqdm(total=opts.total_steps)
    train_dataloader = DataLoader(
        dataset,
        sampler = train_sampler,
        batch_size = opts.per_gpu_batch_size,
        drop_last = True,
        num_workers = 10,
        collate_fn = collator,
    )

    loss, curr_loss = 0.0, 0.0
    epoch = 1
    model.train()
    temp_step = 0
    while step < opts.total_steps:
        epoch += 1
        for i, batch in enumerate(train_dataloader):
            temp_step += 1
            batch = {k:v.cuda() if type(v) != list else v for k, v in batch.items()}
            train_loss, scores, gold_scores = model(**batch)

            train_loss.backward()
            if temp_step % opts.accumulation_steps == 0:
                step += 1
                temp_step = 0
                torch.nn.utils.clip_grad_norm_(model.parameters(), opts.clip)
                optimizer.step()
                scheduler.step()
                model.zero_grad()
                if opts.is_main:
                    bar.update(1)

            train_loss = average_main(train_loss, opts)
            curr_loss += train_loss.item()
            if opts.is_main and step % opts.save_freq == 0:
                save_checkpoint(model, optimizer, scheduler, step, opts, checkpoint_path, f"step-{step}")
            if step > opts.total_steps:
                save_checkpoint(model, optimizer, scheduler, step, opts, checkpoint_path, f"step-{step}")
                break
    
parser = argparse.ArgumentParser()

parser.add_argument("--train_data", required = True, help="training data")
parser.add_argument("--do_train", action='store_true', help="perform training")
parser.add_argument("--scores_path", required=True, help="address to pre-computed profile item score")

parser.add_argument("--max_length_query", type = int, default = 512, help="max length query")
parser.add_argument("--max_length_document", type = int, default = 512, help="max length document")

parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment')
parser.add_argument('--checkpoint_dir', type=str, default='./checkpoint/', help='models are saved here')

parser.add_argument("--per_gpu_batch_size", default=1, type=int, 
           help="Batch size per GPU/CPU for training.")
parser.add_argument("--local-rank", type=int, default=-1,
           help="For distributed training: local_rank")
parser.add_argument("--main_port", type=int, default=-1,
           help="Main port (for multi-node SLURM jobs)")
parser.add_argument('--seed', type=int, default=0, help="random seed for initialization")
parser.add_argument('--save_freq', type=int, default=5000,
           help='save model every <save_freq> steps during training')
parser.add_argument('--warmup_steps', type=int, default=1000, help="number of warmup steps")
parser.add_argument('--total_steps', type=int, default=1000, help="number of training steps")
parser.add_argument('--scheduler_steps', type=int, default=None, 
           help='total number of step for the scheduler, if None then scheduler_total_step = total_step')
parser.add_argument('--accumulation_steps', type=int, default=1, help="number of gradient accumulation steps")
parser.add_argument('--dropout', type=float, default=0.1, help='dropout rate')
parser.add_argument('--lr', type=float, default=1e-5, help='learning rate')
parser.add_argument('--clip', type=float, default=1., help='gradient clipping')
parser.add_argument('--optim', type=str, default='adam', help="optimizer which is used for training")
parser.add_argument('--scheduler', type=str, default='fixed', help="scheduler which is used for training")
parser.add_argument('--weight_decay', type=float, default=0.0, help="weight decay rate")
parser.add_argument('--fixed_lr', action='store_true', help="use a fixed lr")

parser.add_argument('--ctx_size', type=int, default=20, help="number of docs per query for training")

parser.add_argument("--task", required = True, help="task name")
parser.add_argument("--model_path", default="", help="address to a checkpoint to be load")

parser.add_argument('--temperature', type=float, default=1.0, help="temperature for distillation")
parser.add_argument('--cache_dir', default="cache")






if __name__ == "__main__":
    opts = parser.parse_args()

    torch.manual_seed(opts.seed)
    init_distributed_mode(opts)
    init_signal_handler()

    checkpoint_path = Path(opts.checkpoint_dir)/opts.name
    checkpoint_exists = checkpoint_path.exists()
    
    if opts.is_distributed:
        torch.distributed.barrier()
    
    checkpoint_path.mkdir(parents = True, exist_ok = True)
    opts.output_dir = checkpoint_path


    logger = init_logger(
        opts.is_main,
        opts.is_distributed,
        checkpoint_path / 'run.log'
    )

    logger.info(opts)

    model = Contriever.from_pretrained('facebook/contriever', cache_dir = opts.cache_dir)
    tokenizer = AutoTokenizer.from_pretrained('facebook/contriever', cache_dir = opts.cache_dir)
    collator = ReaderToRetreieverCollator(tokenizer = tokenizer, query_max_lenght = opts.max_length_query, document_max_length = opts.max_length_document, number_of_ctx = opts.ctx_size, scores_addr = opts.scores_path)
    
    task = opts.task
    
    reader_tokenizer = AutoTokenizer.from_pretrained('google/flan-t5-base', cache_dir = opts.cache_dir)
    query_corpus_generator = create_query_corpus_generator(task)

    greater_is_better = True

    if task == "LaMP-1":
        train_dataset, labels = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator), get_all_labels(task)
        best_metric_generation = "accuracy"
    elif task == "LaMP-2":
        train_dataset, labels = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator), get_all_labels(task)
        best_metric_generation = "accuracy"
    elif task == "LaMP-3":
        train_dataset, labels = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator), get_all_labels(task)
        best_metric_generation = "mae"
        greater_is_better = False
    elif task == "LaMP-4":
        train_dataset = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator)
        best_metric_generation = "rouge-1"
    elif task == "LaMP-5":
        train_dataset = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator)
        best_metric_generation = "rouge-1"
    elif task == "LaMP-7":
        train_dataset = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator)
        best_metric_generation = "rouge-1"
    elif task == "LaMP-6":
        train_dataset = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator)
        best_metric_generation = "rouge-1"
    
    opts.greater_is_better = greater_is_better
    opts.reader_gold_metric = best_metric_generation

    if not checkpoint_exists and not opts.model_path:
        model = KDReaderToRetrieverTrainer(model = model, args = opts)
        model = model.to(opts.local_rank)
        optimizer, scheduler = set_optim(opts, model)
        step = 0
    elif checkpoint_exists and opts.model_path and opts.do_train:
        model, optimizer, scheduler, opt_checkpoint, step = load_checkpoint(Contriever, opts.model_path, opts)
        model = KDReaderToRetrieverTrainer(model = model, args = opts)
    
    if opts.is_distributed:
        model = torch.nn.parallel.DistributedDataParallel(
            model,
            device_ids=[opts.local_rank],
            output_device=opts.local_rank,
            find_unused_parameters=True,
        )

    if opts.do_train:
        train(opts, model, optimizer, scheduler, step, train_dataset, collator, checkpoint_path)

================================================
FILE: ROPG/train_rl.py
================================================
from pathlib import Path
from utils.distributed import init_distributed_mode, init_signal_handler
import torch
import numpy as np
import os
from torch.utils.data import DataLoader, DistributedSampler, SequentialSampler, RandomSampler
import tqdm
import argparse
from transformers import AutoModel, AutoTokenizer, AutoModelForSeq2SeqLM
from data.collators import ReaderToRetreieverCollator
from data.datasets import ReaderToRetrieverDataset, get_all_labels
from prompts.prompts import create_prompt_generator, create_query_corpus_generator
from trainers.trainer import RLReaderToRetrieverTrainer
from models.retriever import Contriever
from utils.util import average_main
from utils.log import init_logger
from models.optim import set_optim
from utils.util import save_checkpoint, load_checkpoint


def train(opts, model, optimizer, scheduler, step, dataset, collator, checkpoint_path):
    
    if opts.is_main:
        try:
            tb_logger = torch.utils.tensorboard.SummaryWriter(Path(opts.checkpoint_dir)/opts.name)
        except:
           tb_logger = None
           logger.warning('Tensorboard is not available.')

    torch.manual_seed(opts.global_rank + opts.seed)
    if opts.is_distributed:
        train_sampler = DistributedSampler(dataset, num_replicas=opts.n_gpu_per_node, rank=opts.local_rank)
    else:
        train_sampler = RandomSampler(dataset)
    bar = tqdm.tqdm(total=opts.total_steps)
    train_dataloader = DataLoader(
        dataset,
        sampler = train_sampler,
        batch_size = opts.per_gpu_batch_size,
        drop_last = True,
        num_workers = 10,
        collate_fn = collator,
    )

    loss, curr_loss = 0.0, 0.0
    epoch = 1
    model.train()
    temp_step = 0
    while step < opts.total_steps:
        epoch += 1
        for i, batch in enumerate(train_dataloader):
            temp_step += 1
            batch = {k:v.cuda() if type(v) != list else v for k, v in batch.items()}
            train_loss, scores, gold_scores = model(**batch)

            train_loss.backward()
            if temp_step % opts.accumulation_steps == 0:
                step += 1
                bar.update(1)
                temp_step = 0
                torch.nn.utils.clip_grad_norm_(model.parameters(), opts.clip)
                optimizer.step()
                scheduler.step()
                model.zero_grad()

            train_loss = average_main(train_loss, opts)
            curr_loss += train_loss.item()
            if opts.is_main and step % opts.save_freq == 0:
                save_checkpoint(model, optimizer, scheduler, step, opts, checkpoint_path, f"step-{step}")
            if step > opts.total_steps:
                save_checkpoint(model, optimizer, scheduler, step, opts, checkpoint_path, f"step-{step}")
                break
    
parser = argparse.ArgumentParser()

parser.add_argument("--train_data", required = True, help="training data")
parser.add_argument("--do_train", action='store_true', help="perform training")
parser.add_argument("--scores_path", required=True, help="address to pre-computed profile item score")

parser.add_argument("--max_length_query", type = int, default = 512, help="max length query")
parser.add_argument("--max_length_document", type = int, default = 512, help="max length document")

parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment')
parser.add_argument('--checkpoint_dir', type=str, default='./checkpoint/', help='models are saved here')

parser.add_argument("--per_gpu_batch_size", default=1, type=int, 
           help="Batch size per GPU/CPU for training.")
parser.add_argument("--local-rank", type=int, default=-1,
           help="For distributed training: local_rank")
parser.add_argument("--main_port", type=int, default=-1,
           help="Main port (for multi-node SLURM jobs)")
parser.add_argument('--seed', type=int, default=0, help="random seed for initialization")
parser.add_argument('--save_freq', type=int, default=5000,
           help='save model every <save_freq> steps during training')
parser.add_argument('--warmup_steps', type=int, default=1000, help="number of warmup steps")
parser.add_argument('--total_steps', type=int, default=1000, help="number of training steps")
parser.add_argument('--scheduler_steps', type=int, default=None, 
           help='total number of step for the scheduler, if None then scheduler_total_step = total_step')
parser.add_argument('--accumulation_steps', type=int, default=1, help="number of gradient accumulation steps")
parser.add_argument('--dropout', type=float, default=0.1, help='dropout rate')
parser.add_argument('--lr', type=float, default=1e-5, help='learning rate')
parser.add_argument('--clip', type=float, default=1., help='gradient clipping')
parser.add_argument('--optim', type=str, default='adam', help="optimizer which is used for training")
parser.add_argument('--scheduler', type=str, default='fixed', help="scheduler which is used for training")
parser.add_argument('--weight_decay', type=float, default=0.0, help="weight decay rate")
parser.add_argument('--fixed_lr', action='store_true', help="use a fixed lr")

parser.add_argument('--ctx_size', type=int, default=20, help="number of docs per query for training")

parser.add_argument("--task", required = True, help="task name")
parser.add_argument("--model_path", default="", help="address to a checkpoint to be load")
parser.add_argument('--cache_dir', default="cache")


if __name__ == "__main__":
    opts = parser.parse_args()

    torch.manual_seed(opts.seed)
    init_distributed_mode(opts)
    init_signal_handler()

    checkpoint_path = Path(opts.checkpoint_dir)/opts.name
    checkpoint_exists = checkpoint_path.exists()
    
    if opts.is_distributed:
        torch.distributed.barrier()
    
    checkpoint_path.mkdir(parents = True, exist_ok = True)
    opts.output_dir = checkpoint_path


    logger = init_logger(
        opts.is_main,
        opts.is_distributed,
        checkpoint_path / 'run.log'
    )

    logger.info(opts)

    model = Contriever.from_pretrained('facebook/contriever', cache_dir = opts.cache_dir)
    tokenizer = AutoTokenizer.from_pretrained('facebook/contriever', cache_dir = opts.cache_dir)
    collator = ReaderToRetreieverCollator(tokenizer = tokenizer, query_max_lenght = opts.max_length_query, document_max_length = opts.max_length_document, number_of_ctx = opts.ctx_size, scores_addr = opts.scores_path)
    
    task = opts.task
    
    model = Contriever.from_pretrained('facebook/contriever', cache_dir = opts.cache_dir)
    tokenizer = AutoTokenizer.from_pretrained('facebook/contriever', cache_dir = opts.cache_dir)
    collator = ReaderToRetreieverCollator(tokenizer = tokenizer, query_max_lenght = opts.max_length_query, document_max_length = opts.max_length_document, number_of_ctx = opts.ctx_size, scores_addr = opts.scores_path)
    
    task = opts.task
    
    query_corpus_generator = create_query_corpus_generator(task)

    greater_is_better = True

    if task == "LaMP-1":
        train_dataset, labels = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator), get_all_labels(task)
        best_metric_generation = "accuracy"
    elif task == "LaMP-2":
        train_dataset, labels = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator), get_all_labels(task)
        best_metric_generation = "accuracy"
    elif task == "LaMP-3":
        train_dataset, labels = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator), get_all_labels(task)
        best_metric_generation = "mae"
        greater_is_better = False
    elif task == "LaMP-4":
        train_dataset = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator)
        best_metric_generation = "rouge-1"
    elif task == "LaMP-5":
        train_dataset = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator)
        best_metric_generation = "rouge-1"
    elif task == "LaMP-7":
        train_dataset = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator)
        best_metric_generation = "rouge-1"
    elif task == "LaMP-6":
        train_dataset = ReaderToRetrieverDataset(opts.train_data, task, query_corpus_generator)
        best_metric_generation = "rouge-1"
    
    opts.greater_is_better = greater_is_better
    opts.reader_gold_metric = best_metric_generation

    if not checkpoint_exists and not opts.model_path:
        model = RLReaderToRetrieverTrainer(model = model, args = opts)
        model = model.to(opts.local_rank)
        optimizer, scheduler = set_optim(opts, model)
        step = 0
    elif checkpoint_exists and opts.model_path and opts.do_train:
        model, optimizer, scheduler, opt_checkpoint, step = load_checkpoint(Contriever, opts.model_path, opts)
        model = RLReaderToRetrieverTrainer(model = model, args = opts)
    
    
    if opts.is_distributed:
        model = torch.nn.parallel.DistributedDataParallel(
            model,
            device_ids=[opts.local_rank],
            output_device=opts.local_rank,
            find_unused_parameters=True,
        )

    if opts.do_train:
        train(opts, model, optimizer, scheduler, step, train_dataset, collator, checkpoint_path)

================================================
FILE: ROPG/trainers/trainer.py
================================================
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from torch.nn import DataParallel
import torch
from torch import nn
from torch.utils.data import Dataset
from transformers import Trainer, TrainingArguments
import math


def select_elements(tensor, k):
    B = tensor.size(0)
    selected_rows = []
    for i in range(B):
        start_idx = i * k
        end_idx = (i + 1) * k
        selected_rows.append(tensor[i, start_idx:end_idx])
    selected_tensor = torch.stack(selected_rows)
    return selected_tensor

def loss_fn_reinforce(probs, rewards):
    x = torch.mul(probs, rewards)
    x = torch.sum(x) / (probs.shape[0] * probs.shape[1])
    return -x

class KDReaderToRetrieverTrainer(nn.Module):

    def __init__(self, model, args) -> None:
        super().__init__()
        self.model = model
        self.args = args
        
    def forward(
            self, 
            query_input_ids = None,
            query_token_type_ids = None,
            query_attention_mask = None,
            documents_input_ids = None,
            documents_token_type_ids = None,
            documents_attention_mask = None,
            documents_ctxs = None,
            scores_gold = None,
            target_txt = None,
            **kws
        ):
        return self._forward(
            query_input_ids,
            query_token_type_ids,
            query_attention_mask,
            documents_input_ids,
            documents_token_type_ids,
            documents_attention_mask,
            documents_ctxs,
            target_txt,
            scores_gold
        )
    
    def _forward(
            self, 
            query_input_ids,
            query_token_type_ids,
            query_attention_mask,
            documents_input_ids,
            documents_token_type_ids,
            documents_attention_mask,
            documents_ctxs,
            target_txt,
            scores_gold,
            **kws
        ):
        B = documents_token_type_ids.shape[0]
        ctx_size = documents_token_type_ids.shape[1]

        query_reps = self.model(input_ids = query_input_ids, token_type_ids = query_token_type_ids, attention_mask = query_attention_mask)
        docs_reps = self.model(input_ids = documents_input_ids.view(B * ctx_size, -1), token_type_ids = documents_token_type_ids.view(B * ctx_size, -1), attention_mask = documents_attention_mask.view(B * ctx_size, -1))
        docs_reps = docs_reps.view(B, ctx_size, -1)
        scores = torch.einsum("ij,ikj->ik", query_reps, docs_reps)
        probs = torch.softmax(scores / self.args.temperature, dim = -1)
        
        # gold label creation
        if self.args.greater_is_better:
            gold_scores = torch.zeros_like(scores, device = scores.device)
            for i, sample in enumerate(scores_gold):
                for j, score in enumerate(sample):
                    gold_scores[i, j] = (score[self.args.reader_gold_metric])
        else:
            target_numerical = [int(x) for x in target_txt]
            worst_score_array = [max(abs(x-1), abs(x-5)) for x in target_numerical]
            gold_scores = torch.tensor([[float(x) for y in range(ctx_size)] for x in worst_score_array], device = scores.device)
            for i, sample in enumerate(scores_gold):
                for j, score in enumerate(sample):
                    gold_scores[i, j] = (worst_score_array[i] - score[self.args.reader_gold_metric]) / worst_score_array[i]
        
        loss_fn = nn.CrossEntropyLoss()
        loss = loss_fn(probs, torch.softmax(gold_scores  / self.args.temperature, dim = -1))
        
        
        return loss, torch.softmax(scores, dim = -1), gold_scores 
    
class RLReaderToRetrieverTrainer(nn.Module):

    def __init__(self, model, args) -> None:
        super().__init__()
        self.model = model
        self.args = args
        
    def forward(
            self, 
            query_input_ids = None,
            query_token_type_ids = None,
            query_attention_mask = None,
            documents_input_ids = None,
            documents_token_type_ids = None,
            documents_attention_mask = None,
            documents_ctxs = None,
            scores_gold = None,
            target_txt = None,
            **kws
        ):
        return self._forward(
            query_input_ids,
            query_token_type_ids,
            query_attention_mask,
            documents_input_ids,
            documents_token_type_ids,
            documents_attention_mask,
            documents_ctxs,
            target_txt,
            scores_gold
        )
        
    def _forward(
            self, 
            query_input_ids,
            query_token_type_ids,
            query_attention_mask,
            documents_input_ids,
            documents_token_type_ids,
            documents_attention_mask,
            documents_ctxs,
            target_txt,
            scores_gold,
            **kws
        ):
        B = documents_token_type_ids.shape[0]
        ctx_size = documents_token_type_ids.shape[1]

        query_reps = self.model(input_ids = query_input_ids, token_type_ids = query_token_type_ids, attention_mask = query_attention_mask)
        docs_reps = self.model(input_ids = documents_input_ids.view(B * ctx_size, -1), token_type_ids = documents_token_type_ids.view(B * ctx_size, -1), attention_mask = documents_attention_mask.view(B * ctx_size, -1))
        docs_reps = docs_reps.view(B, ctx_size, -1)
        scores = torch.einsum("ij,ikj->ik", query_reps, docs_reps)
        probs = torch.softmax(scores, dim = -1)
        
        sample_idx = probs.multinomial(1)
        
        # gold label creation
        if self.args.greater_is_better:
            gold_scores = torch.zeros_like(scores, device = scores.device)
            for i, sample in enumerate(scores_gold):
                for j, score in enumerate(sample):
                    gold_scores[i, j] = (score[self.args.reader_gold_metric] - sample[0][self.args.reader_gold_metric])
        else:
            target_numerical = [int(x) for x in target_txt]
            worst_score_array = [max(abs(x-1), abs(x-5)) for x in target_numerical]
            gold_scores = torch.tensor([[float(x) for y in range(ctx_size)] for x in worst_score_array], device = scores.device)
            for i, sample in enumerate(scores_gold):
                for j, score in enumerate(sample):
                    gold_scores[i, j] = (sample[0][self.args.reader_gold_metric] - score[self.args.reader_gold_metric]) / worst_score_array[i]
        
        probs = torch.gather(probs, 1, sample_idx)
        gold_scores = torch.gather(gold_scores, 1, sample_idx)
        loss = loss_fn_reinforce(torch.log(probs), gold_scores)
        
        return loss, torch.softmax(scores, dim = -1), gold_scores 

================================================
FILE: ROPG/utils/distributed.py
================================================
from logging import getLogger
import os
import sys
import torch
import socket
import signal
import subprocess
import datetime


logger = getLogger()

def sig_handler(signum, frame):
    logger.warning("Signal handler called with signal " + str(signum))
    prod_id = int(os.environ['SLURM_PROCID'])
    logger.warning("Host: %s - Global rank: %i" % (socket.gethostname(), prod_id))
    if prod_id == 0:
        logger.warning("Requeuing job " + os.environ['SLURM_JOB_ID'])
        os.system('scontrol requeue ' + os.environ['SLURM_JOB_ID'])
    else:
        logger.warning("Not the main process, no need to requeue.")
    sys.exit(-1)


def term_handler(signum, frame):
    logger.warning("Signal handler called with signal " + str(signum))
    logger.warning("Bypassing SIGTERM.")


def init_signal_handler():
    signal.signal(signal.SIGUSR1, sig_handler)
    signal.signal(signal.SIGTERM, term_handler)


def init_distributed_mode(params):
    
    has_local_rank = hasattr(params, 'local_rank')
    if has_local_rank:
        params.local_rank = params.local_rank
    
    if has_local_rank and params.local_rank != -1:

        assert params.main_port == -1

        # read environment variables
        params.global_rank = int(os.environ['RANK'])
        params.world_size = int(os.environ['WORLD_SIZE'])
        params.n_gpu_per_node = int(os.environ['NGPU'])

        # number of nodes / node ID
        params.n_nodes = params.world_size // params.n_gpu_per_node
        params.node_id = params.global_rank // params.n_gpu_per_node
        params.is_distributed = True
    else:
        n_gpu = torch.cuda.device_count()
        params.n_nodes = 1
        params.node_id = 0
        params.local_rank = 0
        params.global_rank = 0
        params.world_size = n_gpu
        params.n_gpu_per_node = n_gpu
        params.is_distributed = False
    
    # define whether this is the master process / if we are in distributed mode
    params.is_main = params.node_id == 0 and params.local_rank == 0
    params.multi_node = params.n_nodes > 1
    params.multi_gpu = params.world_size > 1

    # set GPU device
    if params.is_distributed:
        torch.cuda.set_device(params.local_rank)
        device = torch.device("cuda", params.local_rank)
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    params.device = device

    # initialize multi-GPU
    if params.is_distributed:
        torch.distributed.init_process_group(
            init_method='env://',
            backend='nccl',
            timeout = datetime.timedelta(seconds=36000)
        )

================================================
FILE: ROPG/utils/log.py
================================================
import logging
import torch
import sys

logger = logging.getLogger(__name__)

def init_logger(is_main=True, is_distributed=False, filename=None):
    if is_distributed:
        torch.distributed.barrier()
    handlers = [logging.StreamHandler(sys.stdout)]
    if filename is not None:
        handlers.append(logging.FileHandler(filename = filename))
    logging.basicConfig(
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if is_main else logging.WARN,
        format="[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s",
        handlers=handlers,
    )
    logging.getLogger('transformers.tokenization_utils').setLevel(logging.ERROR)
    logging.getLogger('transformers.tokenization_utils_base').setLevel(logging.ERROR)
    return logger

================================================
FILE: ROPG/utils/util.py
================================================
import os
from logging import getLogger
import torch
from models.optim import set_optim
import torch.distributed as dist
import errno

def load_checkpoint(model_class, dir_path, opt, reset_params=False):
    epoch_path = dir_path
    optimizer_path = os.path.join(epoch_path, "optimizer.pth.tar")
    logger = getLogger()
    logger.info("Loading %s" % epoch_path)
    model = model_class.from_pretrained(epoch_path, local_files_only=True)
    model = model.to(opt.device)
    logger.info("loading checkpoint %s" %optimizer_path)
    checkpoint = torch.load(optimizer_path, map_location=opt.device)
    opt_checkpoint = checkpoint["opt"]
    step = checkpoint["step"]
    if not reset_params:
        optimizer, scheduler = set_optim(opt_checkpoint, model)
        scheduler.load_state_dict(checkpoint["scheduler"])
        optimizer.load_state_dict(checkpoint["optimizer"])
    else:
        optimizer, scheduler = set_optim(opt, model)

    return model, optimizer, scheduler, opt_checkpoint, step

def average_main(x, opt):
    if not opt.is_distributed:
        return x
    if opt.world_size > 1:
        dist.reduce(x, 0, op=dist.ReduceOp.SUM)
        if opt.is_main:
            x = x / opt.world_size
    return x

def symlink_force(target, link_name):
    try:
        os.symlink(target, link_name)
    except OSError as e:
        if e.errno == errno.EEXIST:
            os.remove(link_name)
            os.symlink(target, link_name)
        else:
            raise e

def save_checkpoint(model, optimizer, scheduler, step, opt, dir_path, name):
    model_to_save = model.module.model if hasattr(model, "module") else model.model
    path = os.path.join(dir_path, "checkpoint")
    epoch_path = os.path.join(path, name) #"step-%s" % step)
    os.makedirs(epoch_path, exist_ok=True)
    model_to_save.save_pretrained(epoch_path)
    cp = os.path.join(path, "latest")
    fp = os.path.join(epoch_path, "optimizer.pth.tar")
    checkpoint = {
        "step": step,
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "opt": opt,
    }
    torch.save(checkpoint, fp)
    symlink_force(epoch_path, cp)


================================================
FILE: RSPG/data/collators.py
================================================
import json
import datasets
import torch

class RSPGPreCollator(object):

    def __init__(self, tokenizer, max_length) -> None:
        self.tokenizer = tokenizer
        self.max_len = max_length
        
    def __call__(self, batch):

        inps = [x for ex in batch for x in ex['inputs']]
        inps = self.tokenizer.batch_encode_plus(
            inps,
            max_length=self.max_len,
            padding=True,
            return_tensors='pt',
            truncation=True,
        )
        labels = torch.tensor([x['labels'] for x in batch])
        

        return {
            "id" : [x['id'] for x in batch],
            "input_ids" : inps['input_ids'],
            "attention_mask" : inps['attention_mask'],
            "labels" : labels,
            "outputs" : [x['outputs'] for x in batch],
            "gold" : [x['gold'] for x in batch]
        }

class RSPGPostCollator(object):

    def __init__(self, tokenizer, max_length) -> None:
        self.tokenizer = tokenizer
        self.max_len = max_length
        
    def __call__(self, batch):

        inps = [f"{x} {self.tokenizer.sep_token} {y}" for ex in batch for x, y in zip(ex['inputs'], ex['outputs'])]
        inps = self.tokenizer.batch_encode_plus(
            inps,
            max_length=self.max_len,
            padding=True,
            return_tensors='pt',
            truncation=True,
        )
        labels = torch.tensor([x['labels'] for x in batch])
        

        return {
            "id" : [x['id'] for x in batch],
            "input_ids" : inps['input_ids'],
            "attention_mask" : inps['attention_mask'],
            "labels" : labels,
            "outputs" : [x['outputs'] for x in batch],
            "gold" : [x['gold'] for x in batch]
        }

================================================
FILE: RSPG/data/dataset.py
================================================
from torch.utils.data import Dataset
import json
import datasets
import numpy as np
import random

class RSPGDataset(Dataset):

    def __init__(self, data_addr, smaller_is_better = False) -> None:
        super().__init__()
        with open(data_addr) as file:
            data = json.load(file)
        self.data = data
        self.smaller_is_better = smaller_is_better

    def __getitem__(self, index):
        data = self.data[index]
        did = data['id']
        if self.smaller_is_better:
            worst_score = max(abs(int(data['gold'])-1), abs(int(data['gold'])-5))
            labels = [(worst_score - x) / worst_score for x in self.data[index]['labels']]
        else:
            labels = self.data[index]['labels']
        outputs = self.data[index]['outputs']

        return {
            "id" : did,
            "inputs" : [x.lower() for x in data['inputs']],
            "labels" : labels,
            "outputs" : outputs,
            "gold" : data['gold']
        }
    
    def __len__(self):
        return len(self.data)

================================================
FILE: RSPG/metrics/evaluation.py
================================================
import json
import zipfile
import glob
import os
import shutil
import evaluate

def postprocess_text_classification(preds, labels):
    preds = [str(pred).strip() for pred in preds]
    labels = [str(label).strip() for label in labels]
    return preds, labels

def postprocess_text_generation(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [[label.strip()] for label in labels]

    return preds, labels

def create_metric_f1_accuracy(all_labels):
    f1_metric = evaluate.load("f1")
    accuracy_metric = evaluate.load("accuracy")
    def create_mapping(x):
        try:
            return all_labels.index(x)
        except:
            return -1
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text_classification(decoded_preds, decoded_labels)
        decoded_preds = [create_mapping(x) for x in decoded_preds]
        decoded_labels = [create_mapping(x) for x in decoded_labels]
        result_acc = accuracy_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_f1 = f1_metric.compute(predictions=decoded_preds, references=decoded_labels, labels=list(range(len(all_labels))), average = "macro")
        result = {"accuracy" : result_acc["accuracy"], "f1" : result_f1["f1"]}
        return result
    return compute_metrics

def create_metric_f1_accuracy_sigtest(all_labels):
    f1_metric = evaluate.load("f1")
    accuracy_metric = evaluate.load("accuracy")
    def create_mapping(x):
        try:
            return all_labels.index(x)
        except:
            return -1
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text_classification(decoded_preds, decoded_labels)
        decoded_preds = [create_mapping(x) for x in decoded_preds]
        decoded_labels = [create_mapping(x) for x in decoded_labels]
        results_acc = []
        results_f1 = []
        for pred, gold in zip(decoded_preds, decoded_labels):
            result_acc = accuracy_metric.compute(predictions=[pred], references=[gold])
            result_f1 = f1_metric.compute(predictions=[pred], references=[gold], labels=list(range(len(all_labels))), average = "macro", pos_label = gold)
            results_acc.append(result_acc["accuracy"])
            results_f1.append(result_f1["f1"])
        result = {"accuracy" : results_acc, "f1" : results_f1}
        return result
    return compute_metrics

def create_metric_mae_rmse():
    mse_metric = evaluate.load("mse")
    mae_metric = evaluate.load("mae")
    def create_mapping(x, y):
        try:
            return float(x)
        except:
            print(x)
            y = float(y)
            if abs(1 - y) > abs(5 - y):
                return 1.0
            else:
                return 5.0
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text_classification(decoded_preds, decoded_labels)
        decoded_preds = [create_mapping(x,y) for x,y in zip(decoded_preds, decoded_labels)]
        decoded_labels = [create_mapping(x,x) for x in decoded_labels]
        result_mae = mae_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result_rmse = mse_metric.compute(predictions=decoded_preds, references=decoded_labels, squared = False)
        result = {"MAE" : result_mae["mae"], "RMSE" : result_rmse["mse"]}
        return result
    return compute_metrics

def create_metric_mae_rmse_sigtest():
    mse_metric = evaluate.load("mse")
    mae_metric = evaluate.load("mae")
    def create_mapping(x, y):
        try:
            return float(x)
        except:
            print(x)
            y = float(y)
            if abs(1 - y) > abs(5 - y):
                return 1.0
            else:
                return 5.0
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text_classification(decoded_preds, decoded_labels)
        decoded_preds = [create_mapping(x,y) for x,y in zip(decoded_preds, decoded_labels)]
        decoded_labels = [create_mapping(x,x) for x in decoded_labels]
        results_mae = []
        results_rmse = []
        for pred, gold in zip(decoded_preds, decoded_labels):
            result_mae = mae_metric.compute(predictions=[pred], references=[gold])
            result_rmse = mse_metric.compute(predictions=[pred], references=[gold], squared = False)
            results_mae.append(result_mae["mae"])
            results_rmse.append(result_rmse["mse"])
        result = {"MAE" : results_mae, "RMSE" : results_rmse}
        return result
    return compute_metrics

def create_metric_rouge():
    rouge_metric = evaluate.load('rouge')
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text_generation(decoded_preds, decoded_labels)
        result_rouge = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels)
        result = {"rouge-1" : result_rouge["rouge1"], "rouge-L" : result_rouge["rougeL"]}
        return result
    return compute_metrics

def create_metric_rouge_sigtest():
    rouge_metric = evaluate.load('rouge')
    def compute_metrics(decoded_preds, decoded_labels):
        decoded_preds, decoded_labels = postprocess_text_generation(decoded_preds, decoded_labels)
        result_rouge = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels, use_aggregator = False)
        result = {"rouge-1" : result_rouge["rouge1"], "rouge-L" : result_rouge["rougeL"]}
        return result
    return compute_metrics

class LaMPEvaluation(object):
    
    def __init__(self, all_golds_zip_file_addr = None, single_gold_json_file_addr = None, extract_addr = "./tmp") -> None:
        assert all_golds_zip_file_addr or single_gold_json_file_addr, "The golds should be provided for all datasets or at least one."
        assert not (all_golds_zip_file_addr and single_gold_json_file_addr), "The golds should be provided using zip file or json file not both."
        self.tasks_golds = dict()
        self.extract_addr = extract_addr
        self.evaluate_all_is_possible = False
        if all_golds_zip_file_addr:
            os.makedirs(self.extract_addr, exist_ok=True)
            with zipfile.ZipFile(all_golds_zip_file_addr, 'r') as zobj:
                zobj.extractall(path = extract_addr)
            for file_addr in glob.glob(os.path.join(self.extract_addr, "**/*.json"), recursive=True):
                with open(file_addr) as file:
                    task = json.load(file)
                    self.tasks_golds[task['task']] = task['golds']
            self._empty_dir(self.extract_addr)
            self.evaluate_all_is_possible = True
        if single_gold_json_file_addr:
            with open(single_gold_json_file_addr) as file:
                    task = json.load(file)
                    self.tasks_golds[task['task']] = task['golds']
    
    def _empty_dir(self, directory_path):
        for filename in os.listdir(directory_path):
            file_path = os.path.join(directory_path, filename)
            try:
                if os.path.isfile(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(f'Failed to delete {file_path}. Reason: {e}')

    def _get_all_gold_ids(self, task_name):
        return set([sample['id'] for sample in self.tasks_golds[task_name]])
    
    def _get_all_ids(self, input):
        return set([sample['id'] for sample in input])
    
    def evaluate_all(self, predicts_zipfile_addr):
        assert self.evaluate_all_is_possible, "You did not provide golds for all tasks."
        with zipfile.ZipFile(predicts_zipfile_addr, 'r') as zobj:
            zobj.extractall(path = self.extract_addr)
        results_raw = dict()
        all_task_names = set()
        for file_addr in glob.glob(os.path.join(self.extract_addr, "**/*.json"), recursive=True):
            with open(file_addr) as file:
                preds = json.load(file)
            all_task_names.add(preds['task'])
            results_raw[preds['task']] = self._evaluate_task(preds['golds'], preds['task'])
        self._empty_dir(self.extract_addr)
        assert len(all_task_names) == 7, "The provided results do not cover all the tasks in the benchmark."
        return results_raw

    def evaluate_task(self, predicts_json_addr, task_name):
        with open(predicts_json_addr) as file:
            preds = json.load(file)
        assert preds['task'] == task_name or preds['task'].replace("-","_") == task_name, "The provided task_name and the results do not match."
        assert preds['task'] in self.tasks_golds.keys() or preds['task'].replace("-","_") in self.tasks_golds.keys(), "The provided golds cannot be used to evaluate this task."
        return self._evaluate_task(preds['golds'], task_name)
        
    def _evaluate_task(self, predictions, task_name):
        golds_dict = {y['id']:y['output'] for y in self.tasks_golds[task_name]}
        preds_dict = {x['id']:x['output'] for x in predictions}
        
        gold_ids = self._get_all_gold_ids(task_name)
        pred_ids = self._get_all_ids(predictions)
        print(gold_ids - pred_ids)
        assert gold_ids == pred_ids, "Predictions ids and gold ids do not match."

        if task_name in ["LaMP_1", "LaMP_2"]:
            metric = create_metric_f1_accuracy(self._get_labels(task_name))
        elif task_name == "LaMP_3":
            metric = create_metric_mae_rmse()
        else:
            metric = create_metric_rouge()
        
        gold_ids = list(gold_ids)
        golds = [golds_dict[id] for id in gold_ids]
        preds = [preds_dict[id] for id in gold_ids]
        return metric(preds, golds)
    
    def _evaluate_task_per_sample(self, predictions, task_name):
        golds_dict = {y['id']:y['output'] for y in self.tasks_golds[task_name]}
        preds_dict = {x['id']:x['output'] for x in predictions}
        
        gold_ids = self._get_all_gold_ids(task_name)
        pred_ids = self._get_all_ids(predictions)
        print(gold_ids - pred_ids)
        assert gold_ids == pred_ids, "Predictions ids and gold ids do not match."

        if task_name in ["LaMP_1", "LaMP_2"]:
            metric = create_metric_f1_accuracy_sigtest(self._get_labels(task_name))
        elif task_name == "LaMP_3":
            metric = create_metric_mae_rmse_sigtest()
        else:
            metric = create_metric_rouge_sigtest()
        
        gold_ids = list(gold_ids)
        golds = [golds_dict[id] for id in gold_ids]
        preds = [preds_dict[id] for id in gold_ids]
        return metric(preds, golds)
    
    def _get_labels(self, task_name):
        if task_name == "LaMP_1":
            return ["[1]", "[2]"]
        elif task_name == "LaMP_2":
            return ['sci-fi', 'based on a book', 'comedy', 'action', 'twist ending', 'dystopia', 'dark comedy', 'classic', 'psychology', 'fantasy', 'romance', 'thought-provoking', 'social commentary', 'violence', 'true story']
        elif task_name == "LaMP_3":
            return ["1", "2", "3", "4", "5"]
        else:
            raise ValueError("Invalid task_name")

================================================
FILE: RSPG/modeling/__init__.py
================================================


================================================
FILE: RSPG/modeling/modeling.py
================================================
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import types
import torch
from transformers import PreTrainedModel, AutoModel
import torch.nn.functional as F
from torch import nn
from torch.nn import CrossEntropyLoss
import numpy as np

class RSPG(PreTrainedModel):
    def __init__(self, config, **kwargs):
        super().__init__(config, **kwargs)
        self.model = AutoModel.from_pretrained(config.init_model)
        self.classifier = nn.Linear(self.model.config.hidden_size, 1)
    
    def forward(self, input_ids = None, attention_mask = None, token_type_ids = None, **kwargs):
        output = self.model(
            input_ids = input_ids,
            attention_mask = attention_mask,
            token_type_ids = token_type_ids
        )
        return self.classifier(output.pooler_output).view(-1, self.config.num_labels)

class Trainer(nn.Module):

    def __init__(self, model, temperature = 1.0):
        super().__init__()
        self.model = model
        self.loss_fn_kl = nn.KLDivLoss(reduction="batchmean")
        self.temperature = temperature
    
    def forward(
            self,
            input_ids = None,
            attention_mask = None,
            token_type_ids = None,
            labels = None
        ):

        outputs = self.model(
            input_ids = input_ids,
            attention_mask = attention_mask,
            token_type_ids = token_type_ids
        )

        loss = self.loss_fn_kl(nn.functional.log_softmax(outputs, dim = -1) / self.temperature, nn.functional.softmax(labels, dim = -1))
        
        return loss, outputs

================================================
FILE: RSPG/modeling/optim.py
================================================
import torch

class WarmupLinearScheduler(torch.optim.lr_scheduler.LambdaLR):
    def __init__(self, optimizer, warmup_steps, scheduler_steps, min_ratio, fixed_lr, last_epoch=-1):
        self.warmup_steps = warmup_steps
        self.scheduler_steps = scheduler_steps
        self.min_ratio = min_ratio
        self.fixed_lr = fixed_lr
        super(WarmupLinearScheduler, self).__init__(
            optimizer, self.lr_lambda, last_epoch=last_epoch
        )

    def lr_lambda(self, step):
        if step < self.warmup_steps:
            return (1 - self.min_ratio)*step/float(max(1, self.warmup_steps)) + self.min_ratio

        if self.fixed_lr:
            return 1.0

        return max(0.0,
            1.0 + (self.min_ratio - 1) * (step - self.warmup_steps)/float(max(1.0, self.scheduler_steps - self.warmup_steps)),
        )


class FixedScheduler(torch.optim.lr_scheduler.LambdaLR):
    
    def __init__(self, optimizer, last_epoch=-1):
        super(FixedScheduler, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)
   
    def lr_lambda(self, step):
        return 1.0


def set_optim(opt, model):
    if opt.optim == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr)
    elif opt.optim == 'adamw':
        optimizer = torch.optim.AdamW(model.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)
    if opt.scheduler == 'fixed':
        scheduler = FixedScheduler(optimizer)
    elif opt.scheduler == 'linear':
        if opt.scheduler_steps is None:
            scheduler_steps = opt.total_steps
        else:
            scheduler_steps = opt.scheduler_steps
        scheduler = WarmupLinearScheduler(optimizer, warmup_steps=opt.warmup_steps, scheduler_steps=scheduler_steps, min_ratio=0., fixed_lr=opt.fixed_lr)
    return optimizer, scheduler

================================================
FILE: RSPG/modeling/utils.py
================================================
import os
from logging import getLogger
import torch
from modeling.optim import set_optim
import torch.distributed as dist
import errno

def load_checkpoint(model_class, dir_path, opt, reset_params=False):
    epoch_path = os.path.realpath(dir_path)
    optimizer_path = os.path.join(epoch_path, "optimizer.pth.tar")
    logger = getLogger()
    logger.info("Loading %s" % epoch_path)
    model = model_class.from_pretrained(epoch_path, local_files_only=True)
    model = model.to(opt.device)
    logger.info("loading checkpoint %s" %optimizer_path)
    checkpoint = torch.load(optimizer_path, map_location=opt.device)
    opt_checkpoint = checkpoint["opt"]
    step = checkpoint["step"]
    if not reset_params:
        optimizer, scheduler = set_optim(opt_checkpoint, model)
        scheduler.load_state_dict(checkpoint["scheduler"])
        optimizer.load_state_dict(checkpoint["optimizer"])
    else:
        optimizer, scheduler = set_optim(opt, model)

    return model, optimizer, scheduler, opt_checkpoint, step

def average_main(x, opt):
    if not opt.is_distributed:
        return x
    if opt.world_size > 1:
        dist.reduce(x, 0, op=dist.ReduceOp.SUM)
        if opt.is_main:
            x = x / opt.world_size
    return x

def symlink_force(target, link_name):
    try:
        os.symlink(target, link_name)
    except OSError as e:
        if e.errno == errno.EEXIST:
            os.remove(link_name)
            os.symlink(target, link_name)
        else:
            raise e

def save_checkpoint(model, optimizer, scheduler, step, opt, dir_path, name):
    model_to_save = model.module if hasattr(model, "module") else model
    path = os.path.join(dir_path, "checkpoint")
    epoch_path = os.path.join(path, name) #"step-%s" % step)
    os.makedirs(epoch_path, exist_ok=True)
    model_to_save.save_pretrained(epoch_path)
    cp = os.path.join(path, "latest")
    fp = os.path.join(epoch_path, "optimizer.pth.tar")
    checkpoint = {
        "step": step,
        "optimizer": optimizer.state_dict(),
        "scheduler"
Download .txt
gitextract_ofn4mhve/

├── CC-BY-NC-SA-4.0.txt
├── LaMP/
│   ├── data/
│   │   └── datasets.py
│   ├── evaluate_llm.py
│   ├── metrics/
│   │   ├── classification_metrics.py
│   │   └── generation_metrics.py
│   ├── profile_item_utilization_scorer.py
│   ├── prompts/
│   │   ├── contriever_retriever.py
│   │   ├── prompts.py
│   │   └── utils.py
│   ├── rank_profiles.py
│   ├── requirements.txt
│   ├── retriever_utilization_scorer.py
│   ├── train_llm.py
│   └── utils/
│       └── merge_with_rank.py
├── PEFT/
│   ├── data/
│   │   └── datasets.py
│   ├── evaluate_llm.py
│   ├── requirements.txt
│   └── train_peft.py
├── README.md
├── ROPG/
│   ├── data/
│   │   ├── collators.py
│   │   └── datasets.py
│   ├── models/
│   │   ├── optim.py
│   │   └── retriever.py
│   ├── prompts/
│   │   ├── contriever_retriever.py
│   │   ├── prompts.py
│   │   └── utils.py
│   ├── requirements.txt
│   ├── train_kd.py
│   ├── train_rl.py
│   ├── trainers/
│   │   └── trainer.py
│   └── utils/
│       ├── distributed.py
│       ├── log.py
│       └── util.py
├── RSPG/
│   ├── data/
│   │   ├── collators.py
│   │   └── dataset.py
│   ├── metrics/
│   │   └── evaluation.py
│   ├── modeling/
│   │   ├── __init__.py
│   │   ├── modeling.py
│   │   ├── optim.py
│   │   └── utils.py
│   ├── requirements.txt
│   ├── rspg.py
│   └── utils/
│       ├── __init__.py
│       ├── create_data.py
│       ├── distributed.py
│       └── log.py
├── data/
│   └── avocado/
│       └── create_avocado_dataset.py
└── eval/
    ├── eval_all.py
    ├── eval_task.py
    └── evaluation.py
Download .txt
SYMBOL INDEX (227 symbols across 35 files)

FILE: LaMP/data/datasets.py
  function get_all_labels (line 6) | def get_all_labels(task):
  function create_preprocessor (line 22) | def create_preprocessor(tokenizer, max_length):
  function create_preprocessor_scores (line 30) | def create_preprocessor_scores(tokenizer, max_length):
  function create_preprocessor_scores_seq (line 40) | def create_preprocessor_scores_seq(tokenizer, max_length):
  function convert_to_hf_dataset (line 49) | def convert_to_hf_dataset(dataset, cache_dir):
  class GeneralSeq2SeqDataset (line 55) | class GeneralSeq2SeqDataset(Dataset):
    method __init__ (line 57) | def __init__(self, data_addr, use_profile, task, create_prompt = None)...
    method __getitem__ (line 66) | def __getitem__(self, index):
    method __len__ (line 80) | def __len__(self):
  class GeneralSeq2SeqForScoreGenerationDataset (line 83) | class GeneralSeq2SeqForScoreGenerationDataset(Dataset):
    method __init__ (line 85) | def __init__(self, data_addr, use_profile, task, create_prompt = None,...
    method __getitem__ (line 102) | def __getitem__(self, index):
    method __len__ (line 118) | def __len__(self):

FILE: LaMP/metrics/classification_metrics.py
  function postprocess_text (line 4) | def postprocess_text(preds, labels):
  function create_metric_f1_accuracy (line 10) | def create_metric_f1_accuracy(tokenizer, all_labels):
  function create_metric_f1_accuracy_bert (line 35) | def create_metric_f1_accuracy_bert(all_labels):
  function create_metric_mae_rmse_bert (line 47) | def create_metric_mae_rmse_bert(all_labels):
  function create_metric_mae_rmse (line 59) | def create_metric_mae_rmse(tokenizer, all_labels):
  function create_metric_f1_accuracy_chatgpt (line 89) | def create_metric_f1_accuracy_chatgpt(all_labels):
  function create_metric_mae_rmse_chatgpt (line 108) | def create_metric_mae_rmse_chatgpt(all_labels):

FILE: LaMP/metrics/generation_metrics.py
  function postprocess_text (line 4) | def postprocess_text(preds, labels):
  function create_metric_bleu_rouge_meteor (line 10) | def create_metric_bleu_rouge_meteor(tokenizer):
  function create_metric_bleu_rouge_meteor_chatgpt (line 29) | def create_metric_bleu_rouge_meteor_chatgpt():

FILE: LaMP/prompts/contriever_retriever.py
  function mean_pooling (line 4) | def mean_pooling(token_embeddings, mask):
  function retrieve_top_k_with_contriever (line 9) | def retrieve_top_k_with_contriever(contriver, tokenizer, corpus, profile...

FILE: LaMP/prompts/prompts.py
  function classification_citation_query_corpus_maker (line 7) | def classification_citation_query_corpus_maker(inp, profile):
  function classification_news_query_corpus_maker (line 13) | def classification_news_query_corpus_maker(inp, profile):
  function classification_movies_query_corpus_maker (line 18) | def classification_movies_query_corpus_maker(inp, profile):
  function classification_review_query_corpus_maker (line 23) | def classification_review_query_corpus_maker(inp, profile):
  function generation_news_query_corpus_maker (line 28) | def generation_news_query_corpus_maker(inp, profile):
  function generation_paper_query_corpus_maker (line 33) | def generation_paper_query_corpus_maker(inp, profile):
  function generation_paper_long_query_corpus_maker (line 38) | def generation_paper_long_query_corpus_maker(inp, profile):
  function parphrase_tweet_query_corpus_maker (line 44) | def parphrase_tweet_query_corpus_maker(inp, profile):
  function generation_avocado_query_corpus_maker (line 49) | def generation_avocado_query_corpus_maker(inp, profile):
  function generation_avocado_long_query_corpus_maker (line 54) | def generation_avocado_long_query_corpus_maker(inp, profile):
  function create_classification_citation_prompt (line 59) | def create_classification_citation_prompt(inp, profile, max_length, toke...
  function create_classification_news_prompt (line 71) | def create_classification_news_prompt(inp, profile, max_length, tokenize...
  function create_classification_movies_prompt (line 84) | def create_classification_movies_prompt(inp, profile, max_length, tokeni...
  function create_classification_review_prompt (line 97) | def create_classification_review_prompt(inp, profile, max_length, tokeni...
  function create_generation_news_prompt (line 110) | def create_generation_news_prompt(inp, profile, max_length, tokenizer):
  function create_generation_paper_prompt (line 123) | def create_generation_paper_prompt(inp, profile, max_length, tokenizer):
  function create_generation_paper_long_prompt (line 136) | def create_generation_paper_long_prompt(inp, profile, max_length, tokeni...
  function create_parphrase_tweet_prompt (line 150) | def create_parphrase_tweet_prompt(inp, profile, max_length, tokenizer):
  function create_generation_avocado_prompt (line 163) | def create_generation_avocado_prompt(inp, profile, max_length, tokenizer):
  function create_generation_avocado_long_prompt (line 176) | def create_generation_avocado_long_prompt(inp, profile, max_length, toke...
  function create_prompt_generator (line 189) | def create_prompt_generator(num_retrieve, ret_type = "bm25", is_ranked =...

FILE: LaMP/prompts/utils.py
  function extract_strings_between_quotes (line 1) | def extract_strings_between_quotes(input_string):
  function extract_after_article (line 18) | def extract_after_article(input_string):
  function extract_after_description (line 24) | def extract_after_description(input_string):
  function extract_after_review (line 31) | def extract_after_review(input_string):
  function extract_after_paper (line 37) | def extract_after_paper(input_string):
  function extract_after_abstract (line 43) | def extract_after_abstract(input_string):
  function extract_after_colon (line 49) | def extract_after_colon(input_string):
  function add_string_after_title (line 56) | def add_string_after_title(original_string, string_to_add):
  function batchify (line 64) | def batchify(lst, batch_size):

FILE: LaMP/rank_profiles.py
  function mean_pooling (line 21) | def mean_pooling(token_embeddings, mask):
  function retrieve_top_k_with_contriver (line 26) | def retrieve_top_k_with_contriver(contriver, tokenizer, corpus, profile,...
  function retrieve_top_k_with_bm25 (line 41) | def retrieve_top_k_with_bm25(corpus, profile, query, k):
  function classification_citation_query_corpus_maker (line 48) | def classification_citation_query_corpus_maker(inp, profile, use_date):
  function classification_review_query_corpus_maker (line 58) | def classification_review_query_corpus_maker(inp, profile, use_date):
  function generation_news_query_corpus_maker (line 67) | def generation_news_query_corpus_maker(inp, profile, use_date):
  function generation_paper_query_corpus_maker (line 76) | def generation_paper_query_corpus_maker(inp, profile, use_date):
  function parphrase_tweet_query_corpus_maker (line 85) | def parphrase_tweet_query_corpus_maker(inp, profile, use_date):
  function generation_avocado_query_corpus_maker (line 94) | def generation_avocado_query_corpus_maker(inp, profile, use_date):
  function classification_movies_query_corpus_maker (line 103) | def classification_movies_query_corpus_maker(inp, profile, use_date):

FILE: LaMP/utils/merge_with_rank.py
  function merge (line 4) | def merge(inps, outs, ranks):

FILE: PEFT/data/datasets.py
  function sublists_between_2_and_k (line 8) | def sublists_between_2_and_k(lst, k):
  function sample_sublists (line 15) | def sample_sublists(lst, k, num_samples):
  function get_all_labels (line 30) | def get_all_labels(task):
  function create_preprocessor (line 46) | def create_preprocessor(tokenizer, max_length):
  function convert_to_hf_dataset (line 54) | def convert_to_hf_dataset(dataset, cache_dir):
  function create_input_output_gen_func (line 60) | def create_input_output_gen_func(task):
  function create_per_user_dataset (line 104) | def create_per_user_dataset(data_addr, user_ids, task, cache_dir):
  function create_per_user_dataset_test (line 132) | def create_per_user_dataset_test(data_addr, user_ids, task, cache_dir):
  class GeneralSeq2SeqDataset (line 157) | class GeneralSeq2SeqDataset(Dataset):
    method __init__ (line 159) | def __init__(self, data) -> None:
    method __getitem__ (line 163) | def __getitem__(self, index):
    method __len__ (line 170) | def __len__(self):

FILE: PEFT/train_peft.py
  function is_directory_empty (line 11) | def is_directory_empty(path):

FILE: ROPG/data/collators.py
  class ReaderToRetreieverCollator (line 7) | class ReaderToRetreieverCollator:
    method __init__ (line 9) | def __init__(self, tokenizer, query_max_lenght, document_max_length, n...
    method __call__ (line 19) | def __call__(self, examples: List[dict]):

FILE: ROPG/data/datasets.py
  function get_all_labels (line 5) | def get_all_labels(task):
  function create_preprocessor (line 15) | def create_preprocessor(tokenizer, max_length):
  function create_preprocessor_chatgpt (line 23) | def create_preprocessor_chatgpt(tokenizer, max_length):
  function convert_to_hf_dataset (line 32) | def convert_to_hf_dataset(dataset):
  class GeneralSeq2SeqDataset (line 38) | class GeneralSeq2SeqDataset(Dataset):
    method __init__ (line 40) | def __init__(self, data_addr, use_profile, task, create_prompt = None)...
    method __getitem__ (line 49) | def __getitem__(self, index):
    method __len__ (line 61) | def __len__(self):
  class ReaderToRetrieverDataset (line 64) | class ReaderToRetrieverDataset(Dataset):
    method __init__ (line 66) | def __init__(self, data_addr, task, create_query_corpus, is_llama = Fa...
    method __getitem__ (line 74) | def __getitem__(self, index):
    method __len__ (line 85) | def __len__(self):

FILE: ROPG/models/optim.py
  class WarmupLinearScheduler (line 3) | class WarmupLinearScheduler(torch.optim.lr_scheduler.LambdaLR):
    method __init__ (line 4) | def __init__(self, optimizer, warmup_steps, scheduler_steps, min_ratio...
    method lr_lambda (line 13) | def lr_lambda(self, step):
  class FixedScheduler (line 25) | class FixedScheduler(torch.optim.lr_scheduler.LambdaLR):
    method __init__ (line 27) | def __init__(self, optimizer, last_epoch=-1):
    method lr_lambda (line 30) | def lr_lambda(self, step):
  function set_optim (line 34) | def set_optim(opt, model):

FILE: ROPG/models/retriever.py
  class Contriever (line 6) | class Contriever(BertModel):
    method __init__ (line 7) | def __init__(self, config, pooling="average", **kwargs):
    method forward (line 12) | def forward(

FILE: ROPG/prompts/contriever_retriever.py
  function mean_pooling (line 4) | def mean_pooling(token_embeddings, mask):
  function retrieve_top_k_with_contriever (line 9) | def retrieve_top_k_with_contriever(contriver, tokenizer, corpus, profile...

FILE: ROPG/prompts/prompts.py
  function classification_citation_query_corpus_maker (line 7) | def classification_citation_query_corpus_maker(inp, profile):
  function classification_news_query_corpus_maker (line 13) | def classification_news_query_corpus_maker(inp, profile):
  function classification_movies_query_corpus_maker (line 18) | def classification_movies_query_corpus_maker(inp, profile):
  function classification_review_query_corpus_maker (line 23) | def classification_review_query_corpus_maker(inp, profile):
  function generation_news_query_corpus_maker (line 28) | def generation_news_query_corpus_maker(inp, profile):
  function generation_paper_query_corpus_maker (line 33) | def generation_paper_query_corpus_maker(inp, profile):
  function generation_paper_long_query_corpus_maker (line 38) | def generation_paper_long_query_corpus_maker(inp, profile):
  function parphrase_tweet_query_corpus_maker (line 44) | def parphrase_tweet_query_corpus_maker(inp, profile):
  function generation_avocado_query_corpus_maker (line 49) | def generation_avocado_query_corpus_maker(inp, profile):
  function generation_avocado_long_query_corpus_maker (line 54) | def generation_avocado_long_query_corpus_maker(inp, profile):
  function create_classification_citation_prompt (line 59) | def create_classification_citation_prompt(inp, profile, max_length, toke...
  function create_classification_news_prompt (line 71) | def create_classification_news_prompt(inp, profile, max_length, tokenize...
  function create_classification_review_prompt (line 84) | def create_classification_review_prompt(inp, profile, max_length, tokeni...
  function create_generation_news_prompt (line 97) | def create_generation_news_prompt(inp, profile, max_length, tokenizer):
  function create_generation_paper_prompt (line 110) | def create_generation_paper_prompt(inp, profile, max_length, tokenizer):
  function create_generation_paper_long_prompt (line 123) | def create_generation_paper_long_prompt(inp, profile, max_length, tokeni...
  function create_parphrase_tweet_prompt (line 136) | def create_parphrase_tweet_prompt(inp, profile, max_length, tokenizer):
  function create_generation_avocado_prompt (line 149) | def create_generation_avocado_prompt(inp, profile, max_length, tokenizer):
  function create_generation_avocado_long_prompt (line 162) | def create_generation_avocado_long_prompt(inp, profile, max_length, toke...
  function create_classification_movies_prompt (line 175) | def create_classification_movies_prompt(inp, profile, max_length, tokeni...
  function create_query_corpus_generator (line 188) | def create_query_corpus_generator(task):
  function create_prompt_generator (line 207) | def create_prompt_generator(num_retrieve, ret_type = "bm25", is_ranked =...

FILE: ROPG/prompts/utils.py
  function extract_strings_between_quotes (line 3) | def extract_strings_between_quotes(input_string):
  function extract_after_article (line 8) | def extract_after_article(input_string):
  function extract_after_description (line 14) | def extract_after_description(input_string):
  function extract_after_review (line 20) | def extract_after_review(input_string):
  function extract_after_paper (line 26) | def extract_after_paper(input_string):
  function extract_after_colon (line 32) | def extract_after_colon(input_string):
  function extract_after_abstract (line 38) | def extract_after_abstract(input_string):
  function add_string_after_title (line 45) | def add_string_after_title(original_string, string_to_add):
  function batchify (line 53) | def batchify(lst, batch_size):

FILE: ROPG/train_kd.py
  function train (line 21) | def train(opts, model, optimizer, scheduler, step, dataset, collator, ch...

FILE: ROPG/train_rl.py
  function train (line 21) | def train(opts, model, optimizer, scheduler, step, dataset, collator, ch...

FILE: ROPG/trainers/trainer.py
  function select_elements (line 11) | def select_elements(tensor, k):
  function loss_fn_reinforce (line 21) | def loss_fn_reinforce(probs, rewards):
  class KDReaderToRetrieverTrainer (line 26) | class KDReaderToRetrieverTrainer(nn.Module):
    method __init__ (line 28) | def __init__(self, model, args) -> None:
    method forward (line 33) | def forward(
    method _forward (line 58) | def _forward(
  class RLReaderToRetrieverTrainer (line 100) | class RLReaderToRetrieverTrainer(nn.Module):
    method __init__ (line 102) | def __init__(self, model, args) -> None:
    method forward (line 107) | def forward(
    method _forward (line 132) | def _forward(

FILE: ROPG/utils/distributed.py
  function sig_handler (line 13) | def sig_handler(signum, frame):
  function term_handler (line 25) | def term_handler(signum, frame):
  function init_signal_handler (line 30) | def init_signal_handler():
  function init_distributed_mode (line 35) | def init_distributed_mode(params):

FILE: ROPG/utils/log.py
  function init_logger (line 7) | def init_logger(is_main=True, is_distributed=False, filename=None):

FILE: ROPG/utils/util.py
  function load_checkpoint (line 8) | def load_checkpoint(model_class, dir_path, opt, reset_params=False):
  function average_main (line 28) | def average_main(x, opt):
  function symlink_force (line 37) | def symlink_force(target, link_name):
  function save_checkpoint (line 47) | def save_checkpoint(model, optimizer, scheduler, step, opt, dir_path, na...

FILE: RSPG/data/collators.py
  class RSPGPreCollator (line 5) | class RSPGPreCollator(object):
    method __init__ (line 7) | def __init__(self, tokenizer, max_length) -> None:
    method __call__ (line 11) | def __call__(self, batch):
  class RSPGPostCollator (line 33) | class RSPGPostCollator(object):
    method __init__ (line 35) | def __init__(self, tokenizer, max_length) -> None:
    method __call__ (line 39) | def __call__(self, batch):

FILE: RSPG/data/dataset.py
  class RSPGDataset (line 7) | class RSPGDataset(Dataset):
    method __init__ (line 9) | def __init__(self, data_addr, smaller_is_better = False) -> None:
    method __getitem__ (line 16) | def __getitem__(self, index):
    method __len__ (line 34) | def __len__(self):

FILE: RSPG/metrics/evaluation.py
  function postprocess_text_classification (line 8) | def postprocess_text_classification(preds, labels):
  function postprocess_text_generation (line 13) | def postprocess_text_generation(preds, labels):
  function create_metric_f1_accuracy (line 19) | def create_metric_f1_accuracy(all_labels):
  function create_metric_f1_accuracy_sigtest (line 37) | def create_metric_f1_accuracy_sigtest(all_labels):
  function create_metric_mae_rmse (line 60) | def create_metric_mae_rmse():
  function create_metric_mae_rmse_sigtest (line 83) | def create_metric_mae_rmse_sigtest():
  function create_metric_rouge (line 111) | def create_metric_rouge():
  function create_metric_rouge_sigtest (line 120) | def create_metric_rouge_sigtest():
  class LaMPEvaluation (line 129) | class LaMPEvaluation(object):
    method __init__ (line 131) | def __init__(self, all_golds_zip_file_addr = None, single_gold_json_fi...
    method _empty_dir (line 152) | def _empty_dir(self, directory_path):
    method _get_all_gold_ids (line 163) | def _get_all_gold_ids(self, task_name):
    method _get_all_ids (line 166) | def _get_all_ids(self, input):
    method evaluate_all (line 169) | def evaluate_all(self, predicts_zipfile_addr):
    method evaluate_task (line 184) | def evaluate_task(self, predicts_json_addr, task_name):
    method _evaluate_task (line 191) | def _evaluate_task(self, predictions, task_name):
    method _evaluate_task_per_sample (line 212) | def _evaluate_task_per_sample(self, predictions, task_name):
    method _get_labels (line 233) | def _get_labels(self, task_name):

FILE: RSPG/modeling/modeling.py
  class RSPG (line 15) | class RSPG(PreTrainedModel):
    method __init__ (line 16) | def __init__(self, config, **kwargs):
    method forward (line 21) | def forward(self, input_ids = None, attention_mask = None, token_type_...
  class Trainer (line 29) | class Trainer(nn.Module):
    method __init__ (line 31) | def __init__(self, model, temperature = 1.0):
    method forward (line 37) | def forward(

FILE: RSPG/modeling/optim.py
  class WarmupLinearScheduler (line 3) | class WarmupLinearScheduler(torch.optim.lr_scheduler.LambdaLR):
    method __init__ (line 4) | def __init__(self, optimizer, warmup_steps, scheduler_steps, min_ratio...
    method lr_lambda (line 13) | def lr_lambda(self, step):
  class FixedScheduler (line 25) | class FixedScheduler(torch.optim.lr_scheduler.LambdaLR):
    method __init__ (line 27) | def __init__(self, optimizer, last_epoch=-1):
    method lr_lambda (line 30) | def lr_lambda(self, step):
  function set_optim (line 34) | def set_optim(opt, model):

FILE: RSPG/modeling/utils.py
  function load_checkpoint (line 8) | def load_checkpoint(model_class, dir_path, opt, reset_params=False):
  function average_main (line 28) | def average_main(x, opt):
  function symlink_force (line 37) | def symlink_force(target, link_name):
  function save_checkpoint (line 47) | def save_checkpoint(model, optimizer, scheduler, step, opt, dir_path, na...

FILE: RSPG/rspg.py
  function train (line 67) | def train(opts, model, optimizer, scheduler, step, dataset, collator, ch...
  function evaluate (line 133) | def evaluate(model, dataset, collator, opt, step, logger, evaluator, tes...

FILE: RSPG/utils/create_data.py
  function merge (line 5) | def merge(inps, outs, label_files, input_files, score_name):

FILE: RSPG/utils/distributed.py
  function sig_handler (line 14) | def sig_handler(signum, frame):
  function term_handler (line 26) | def term_handler(signum, frame):
  function init_signal_handler (line 31) | def init_signal_handler():
  function init_distributed_mode (line 36) | def init_distributed_mode(params):

FILE: RSPG/utils/log.py
  function init_logger (line 7) | def init_logger(is_main=True, is_distributed=False, filename=None):

FILE: data/avocado/create_avocado_dataset.py
  function empty_dir (line 10) | def empty_dir(directory_path):
  function process_file (line 23) | def process_file(file_addr):

FILE: eval/evaluation.py
  function postprocess_text_classification (line 8) | def postprocess_text_classification(preds, labels):
  function postprocess_text_generation (line 13) | def postprocess_text_generation(preds, labels):
  function create_metric_f1_accuracy (line 19) | def create_metric_f1_accuracy(all_labels):
  function create_metric_mae_rmse (line 37) | def create_metric_mae_rmse():
  function create_metric_rouge (line 60) | def create_metric_rouge():
  class LaMPEvaluation (line 69) | class LaMPEvaluation(object):
    method __init__ (line 71) | def __init__(self, all_golds_zip_file_addr = None, single_gold_json_fi...
    method _empty_dir (line 92) | def _empty_dir(self, directory_path):
    method _get_all_gold_ids (line 103) | def _get_all_gold_ids(self, task_name):
    method _get_all_ids (line 106) | def _get_all_ids(self, input):
    method evaluate_all (line 109) | def evaluate_all(self, predicts_zipfile_addr):
    method evaluate_task (line 124) | def evaluate_task(self, predicts_json_addr, task_name):
    method _evaluate_task (line 131) | def _evaluate_task(self, predictions, task_name):
    method _get_labels (line 152) | def _get_labels(self, task_name):
Condensed preview — 50 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (237K chars).
[
  {
    "path": "CC-BY-NC-SA-4.0.txt",
    "chars": 20146,
    "preview": "Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International\n\nCreative Commons Corporation (\"Creative Commons"
  },
  {
    "path": "LaMP/data/datasets.py",
    "chars": 4657,
    "preview": "from torch.utils.data import Dataset\nimport json\nimport datasets\nimport torch\n\ndef get_all_labels(task):\n    if task == "
  },
  {
    "path": "LaMP/evaluate_llm.py",
    "chars": 4735,
    "preview": "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, AutoModelForCau"
  },
  {
    "path": "LaMP/metrics/classification_metrics.py",
    "chars": 5826,
    "preview": "import numpy as np\nimport evaluate\n\ndef postprocess_text(preds, labels):\n    preds = [pred.strip() for pred in preds]\n  "
  },
  {
    "path": "LaMP/metrics/generation_metrics.py",
    "chars": 2316,
    "preview": "import numpy as np\nimport evaluate\n\ndef postprocess_text(preds, labels):\n    preds = [pred.strip() for pred in preds]\n  "
  },
  {
    "path": "LaMP/profile_item_utilization_scorer.py",
    "chars": 5091,
    "preview": "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, AutoModelForCau"
  },
  {
    "path": "LaMP/prompts/contriever_retriever.py",
    "chars": 1221,
    "preview": "import torch\nfrom prompts.utils import batchify\n\ndef mean_pooling(token_embeddings, mask):\n    token_embeddings = token_"
  },
  {
    "path": "LaMP/prompts/prompts.py",
    "chars": 14241,
    "preview": "from rank_bm25 import BM25Okapi\nfrom transformers import AutoTokenizer, AutoModel\nfrom prompts.utils import extract_stri"
  },
  {
    "path": "LaMP/prompts/utils.py",
    "chars": 2095,
    "preview": "def extract_strings_between_quotes(input_string):\n    output_list = []\n    inside_quotes = False\n    current_string = ''"
  },
  {
    "path": "LaMP/rank_profiles.py",
    "chars": 6809,
    "preview": "import torch\nfrom prompts.utils import batchify\nfrom transformers import AutoModel, AutoTokenizer\nimport json\nimport tqd"
  },
  {
    "path": "LaMP/requirements.txt",
    "chars": 98,
    "preview": "mail_parser==3.15.0\nnumpy==1.24.2\nrank_bm25==0.2.2\ntorch==2.0.0\ntqdm==4.65.0\ntransformers==4.27.1\n"
  },
  {
    "path": "LaMP/retriever_utilization_scorer.py",
    "chars": 5041,
    "preview": "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, AutoModelForCau"
  },
  {
    "path": "LaMP/train_llm.py",
    "chars": 8752,
    "preview": "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments\nfrom transformer"
  },
  {
    "path": "LaMP/utils/merge_with_rank.py",
    "chars": 1498,
    "preview": "import json \nimport argparse\n\ndef merge(inps, outs, ranks):\n    for inp in inps:\n        for o in outs:\n            if o"
  },
  {
    "path": "PEFT/data/datasets.py",
    "chars": 6460,
    "preview": "from torch.utils.data import Dataset\nimport json\nimport datasets\nimport torch\nimport random\nfrom itertools import combin"
  },
  {
    "path": "PEFT/evaluate_llm.py",
    "chars": 4101,
    "preview": "import argparse\nfrom transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments\n"
  },
  {
    "path": "PEFT/requirements.txt",
    "chars": 155,
    "preview": "datasets==2.8.0\nregex==2022.10.31\nsentencepiece==0.1.97\ntokenizers==0.11.1\ntorch==2.0.1\ntqdm==4.64.1\ntransformers==4.28."
  },
  {
    "path": "PEFT/train_peft.py",
    "chars": 4176,
    "preview": "import argparse\nfrom transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments\n"
  },
  {
    "path": "README.md",
    "chars": 19525,
    "preview": "# Codes for papers on Large Language Models Personalization (LaMP)\n\n[LaMP: When Large Language Models Meet Personalizati"
  },
  {
    "path": "ROPG/data/collators.py",
    "chars": 2990,
    "preview": "from typing import Any, List\nimport numpy as np\nimport torch\nimport json\n\n\nclass ReaderToRetreieverCollator:\n\n    def __"
  },
  {
    "path": "ROPG/data/datasets.py",
    "chars": 3349,
    "preview": "from torch.utils.data import Dataset\nimport json\nimport datasets\n\ndef get_all_labels(task):\n    if task == \"LaMP-1\":\n   "
  },
  {
    "path": "ROPG/models/optim.py",
    "chars": 1805,
    "preview": "import torch\n\nclass WarmupLinearScheduler(torch.optim.lr_scheduler.LambdaLR):\n    def __init__(self, optimizer, warmup_s"
  },
  {
    "path": "ROPG/models/retriever.py",
    "chars": 1689,
    "preview": "from typing import Any\nfrom transformers import BertModel\nfrom transformers.configuration_utils import PretrainedConfig\n"
  },
  {
    "path": "ROPG/prompts/contriever_retriever.py",
    "chars": 1221,
    "preview": "import torch\nfrom prompts.utils import batchify\n\ndef mean_pooling(token_embeddings, mask):\n    token_embeddings = token_"
  },
  {
    "path": "ROPG/prompts/prompts.py",
    "chars": 14630,
    "preview": "from rank_bm25 import BM25Okapi\nfrom transformers import AutoTokenizer, AutoModel\nfrom prompts.utils import extract_stri"
  },
  {
    "path": "ROPG/prompts/utils.py",
    "chars": 1762,
    "preview": "import re\n\ndef extract_strings_between_quotes(input_string):\n    pattern = r'\"(.*?)\"'\n    titles = re.findall(pattern, i"
  },
  {
    "path": "ROPG/requirements.txt",
    "chars": 92,
    "preview": "evaluate==0.4.0\nnumpy==1.24.3\nrank_bm25==0.2.2\ntorch==2.0.1\ntransformers==4.29.2\nrouge_score"
  },
  {
    "path": "ROPG/train_kd.py",
    "chars": 9018,
    "preview": "from pathlib import Path\nfrom utils.distributed import init_distributed_mode, init_signal_handler\nimport torch\nimport nu"
  },
  {
    "path": "ROPG/train_rl.py",
    "chars": 9239,
    "preview": "from pathlib import Path\nfrom utils.distributed import init_distributed_mode, init_signal_handler\nimport torch\nimport nu"
  },
  {
    "path": "ROPG/trainers/trainer.py",
    "chars": 6830,
    "preview": "from dataclasses import dataclass\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\nfrom torch.nn imp"
  },
  {
    "path": "ROPG/utils/distributed.py",
    "chars": 2602,
    "preview": "from logging import getLogger\nimport os\nimport sys\nimport torch\nimport socket\nimport signal\nimport subprocess\nimport dat"
  },
  {
    "path": "ROPG/utils/log.py",
    "chars": 773,
    "preview": "import logging\nimport torch\nimport sys\n\nlogger = logging.getLogger(__name__)\n\ndef init_logger(is_main=True, is_distribut"
  },
  {
    "path": "ROPG/utils/util.py",
    "chars": 2153,
    "preview": "import os\nfrom logging import getLogger\nimport torch\nfrom models.optim import set_optim\nimport torch.distributed as dist"
  },
  {
    "path": "RSPG/data/collators.py",
    "chars": 1766,
    "preview": "import json\nimport datasets\nimport torch\n\nclass RSPGPreCollator(object):\n\n    def __init__(self, tokenizer, max_length) "
  },
  {
    "path": "RSPG/data/dataset.py",
    "chars": 1049,
    "preview": "from torch.utils.data import Dataset\nimport json\nimport datasets\nimport numpy as np\nimport random\n\nclass RSPGDataset(Dat"
  },
  {
    "path": "RSPG/metrics/evaluation.py",
    "chars": 11295,
    "preview": "import json\nimport zipfile\nimport glob\nimport os\nimport shutil\nimport evaluate\n\ndef postprocess_text_classification(pred"
  },
  {
    "path": "RSPG/modeling/__init__.py",
    "chars": 0,
    "preview": ""
  },
  {
    "path": "RSPG/modeling/modeling.py",
    "chars": 1737,
    "preview": "# Copyright (c) Facebook, Inc. and its affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the lic"
  },
  {
    "path": "RSPG/modeling/optim.py",
    "chars": 1805,
    "preview": "import torch\n\nclass WarmupLinearScheduler(torch.optim.lr_scheduler.LambdaLR):\n    def __init__(self, optimizer, warmup_s"
  },
  {
    "path": "RSPG/modeling/utils.py",
    "chars": 2162,
    "preview": "import os\nfrom logging import getLogger\nimport torch\nfrom modeling.optim import set_optim\nimport torch.distributed as di"
  },
  {
    "path": "RSPG/requirements.txt",
    "chars": 151,
    "preview": "datasets==2.8.0\nregex==2022.10.31\nsentencepiece==0.1.97\ntokenizers==0.11.1\ntorch==2.0.1\ntqdm==4.64.1\ntransformers==4.28."
  },
  {
    "path": "RSPG/rspg.py",
    "chars": 11716,
    "preview": "import argparse\nimport torch\nimport json\nimport os\nfrom pathlib import Path\nfrom utils.log import init_logger\nfrom pathl"
  },
  {
    "path": "RSPG/utils/__init__.py",
    "chars": 0,
    "preview": ""
  },
  {
    "path": "RSPG/utils/create_data.py",
    "chars": 1956,
    "preview": "import argparse\nimport json\nimport os\n\ndef merge(inps, outs, label_files, input_files, score_name):\n    for inp in inps:"
  },
  {
    "path": "RSPG/utils/distributed.py",
    "chars": 2615,
    "preview": "from logging import getLogger\nimport os\nimport sys\nimport torch\nimport socket\nimport signal\nimport subprocess\nimport dat"
  },
  {
    "path": "RSPG/utils/log.py",
    "chars": 773,
    "preview": "import logging\nimport torch\nimport sys\n\nlogger = logging.getLogger(__name__)\n\ndef init_logger(is_main=True, is_distribut"
  },
  {
    "path": "data/avocado/create_avocado_dataset.py",
    "chars": 6040,
    "preview": "import zipfile\nimport glob\nimport os\nimport shutil\nimport json\nimport tqdm\nimport mailparser\nimport argparse\n\ndef empty_"
  },
  {
    "path": "eval/eval_all.py",
    "chars": 833,
    "preview": "from evaluation import LaMPEvaluation\nimport argparse\nimport json\n\nparser = argparse.ArgumentParser()\n\nparser.add_argume"
  },
  {
    "path": "eval/eval_task.py",
    "chars": 820,
    "preview": "from evaluation import LaMPEvaluation\nimport argparse\nimport json\n\nparser = argparse.ArgumentParser()\n\nparser.add_argume"
  },
  {
    "path": "eval/evaluation.py",
    "chars": 7432,
    "preview": "import json\nimport zipfile\nimport glob\nimport os\nimport shutil\nimport evaluate\n\ndef postprocess_text_classification(pred"
  }
]

About this extraction

This page contains the full source code of the LaMP-Benchmark/LaMP GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 50 files (221.9 KB), approximately 53.9k tokens, and a symbol index with 227 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.

Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.

Copied to clipboard!