Input:
import os
import math
import json
import numpy as np
import pandas as pd
import torch
import random
import csv
from collections import defaultdict
from transformers import BartTokenizer, T5Tokenizer
from transformers import AdamW, get_linear_schedule_with_warmup
from utils import *
from scoring.fluency_scorer import FluencyScorer
from scoring.saliency_scorer import SaliencyBERTScore
from scoring.simplicity_scorer import SimplicityTextScore
from scoring.guardrails import *
from scoring.aggregate_scorer import ScorerWrapper
from GAP.data_relations_as_nodes import GAPDataloader, EventDataset, WebNLGDataset
from GAP.data_relations_as_nodes import evaluate_bleu, get_t_emb_dim
from tqdm import tqdm, trange
from rake_nltk import Rake
from evaluate import load
from sentence_similarity import sentence_similarity
from GAP.modeling_gap_type import GAPBartForConditionalGeneration as GAP_Type_model
from GAP.modeling_gap import GAPBartForConditionalGeneration as GAP_model
from T5.modeling_t5 import T5ForConditionalGeneration as T5

Context:
simplify_batched.py

# import yake
#T5

bertscore = load("bertscore")
## sentence model for merge
phrase_model = sentence_similarity(model_name='distilbert-base-uncased',embedding_type='cls_token_embedding')

## for sentence checking
ner_check = NERInaccuracyPenalty()

def run(args, logger):
    #load in model for graph-to-text and tokenizer
    checkpoint = args.model_path
    tokenizer_path = args.tokenizer_path
    if 't5' in tokenizer_path:
        tokenizer = T5Tokenizer.from_pretrained(tokenizer_path)
    else:
        tokenizer = BartTokenizer.from_pretrained(tokenizer_path)

    n_gpu = torch.cuda.device_count()
    if n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

    if args.type_encoding:
        t_emb_dim = get_t_emb_dim(args)
        model = GAP_Type_model.from_pretrained(checkpoint,t_emb_dim=t_emb_dim)
    else:
        if 't5' in args.model_path:

GAP/data_relations_as_nodes.py
class WebNLGDataset(Dataset):
    def __init__(self, logger, args, data_path, tokenizer, mode):
        self.data_path = data_path
        self.tokenizer = tokenizer
        self.topology = {"entity-entity": args.entity_entity,
                           "entity-relation": args.entity_relation,
                           "relation-entity": args.relation_entity,
                           "relation-relation": args.relation_relation
                          }

        with open(self.data_path + '.json', 'r') as f:
            self.data = json.load(f)

        print("Total samples = {}".format(len(self.data)))

        assert type(self.data) == list
        assert all(["id" in d for d in self.data]), self.data[0].keys()
        if type(self.data[0]["id"]) == int:
            for i in range(len(self.data)):
                self.data[i]["id"] = str(self.data[i]["id"])

        self.args = args
        self.data_type = mode
        self.metric = "BLEU"

        self.head_ids, self.rel_ids, self.tail_ids = self.tokenizer.encode(' [head]', add_special_tokens=False), \
                                                     self.tokenizer.encode(' [relation]', add_special_tokens=False), \
                                                     self.tokenizer.encode(' [tail]', add_special_tokens=False)

        self.graph_ids, self.text_ids = self.tokenizer.encode(' [graph]', add_special_tokens=False), \
                                        self.tokenizer.encode(' [text]', add_special_tokens=False)

        if self.args.model_name == "bart":
            self.mask_token = self.tokenizer.mask_token
            self.mask_token_id = self.tokenizer.mask_token_id
        else:
            self.mask_token = self.tokenizer.additional_special_tokens[0]
            self.mask_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.additional_special_tokens[0])

        if self.args.model_name == "bart":
            if self.args.append_another_bos:
                self.add_bos_id = [self.tokenizer.bos_token_id] * 2
            else:
                self.add_bos_id = [self.tokenizer.bos_token_id]
        else:
            self.add_bos_id = []

    def __len__(self):
        return len(self.data)

    def linearize_v2(self, entity, entity_change, head_ids, rel_ids, tail_ids,
                        relation_change, cnt_edge, adj_matrix):
        # string_label: encoder ids
        # string_label_tokens: encoder tokens

        if len(entity[0]) == 0:
            return [], '', [], [], cnt_edge, adj_matrix
        nodes, edges = [], []
        string_label = copy.deepcopy(head_ids)
        string_label_tokens = ' [head]'
        nodes.extend([-1] * len(string_label))
        edges.extend([-1] * len(string_label))

        string_label += entity_change[entity[0]][0]
        string_label_tokens += ' {}'.format(entity[0])
        nodes.extend([entity_change[entity[0]][1]] * len(entity_change[entity[0]][0]))
        edges.extend([-1] * len(entity_change[entity[0]][0]))

        for rel in entity[2]:
            if len(rel[0]) != 0 and len(rel[1]) != 0:
                rel_label = relation_change[rel[0]]
                rel_ent_label = entity_change[rel[0]][1]
                rel_label_token = copy.deepcopy(rel[0])
                words_label = rel_ids + rel_label + tail_ids + entity_change[rel[1]][0]
                words_label_tokens = ' [relation] {} [tail] {}'.format(rel_label_token, rel[1])
                nodes.extend(
                        ([-1] * len(rel_ids)) + ([entity_change[rel[0]][1]] * len(rel_label)) + ([-1] * len(tail_ids)) + ([entity_change[rel[1]][1]] * len(
                            entity_change[rel[1]][0])))

                edges.extend([-1] * len(rel_ids) + [cnt_edge] * len(rel_label) + [-1] * (
                            len(tail_ids) + len(entity_change[rel[1]][0])))
                if entity_change[entity[0]][1] < len(adj_matrix) and entity_change[rel[1]][1] < len(adj_matrix):
                    if self.topology['entity-entity']:
                        adj_matrix[entity_change[entity[0]][1]][entity_change[rel[1]][1]] = 1
                        adj_matrix[entity_change[rel[1]][1]][entity_change[entity[0]][1]] = 1

                    if self.topology['entity-relation']:
                        adj_matrix[entity_change[entity[0]][1]][entity_change[rel[0]][1]] = 2
                        adj_matrix[entity_change[rel[1]][1]][entity_change[rel[0]][1]] = 2

                    if self.topology['relation-entity']:
                        adj_matrix[entity_change[rel[0]][1]][entity_change[entity[0]][1]] = 3
                        adj_matrix[entity_change[rel[0]][1]][entity_change[rel[1]][1]] = 3

                    if not self.topology['relation-entity'] and not self.topology['relation-relation']:
                        adj_matrix[entity_change[rel[0]][1]][entity_change[rel[0]][1]] = 10

                    if not self.topology['entity-relation'] and not self.topology['entity-entity']:
                        adj_matrix[entity_change[entity[0]][1]][entity_change[entity[0]][1]] = 10
                        adj_matrix[entity_change[rel[1]][1]][entity_change[rel[1]][1]] = 10

                cnt_edge += 1
                string_label += words_label
                string_label_tokens += words_label_tokens

        assert len(string_label) == len(nodes) == len(edges)

        return string_label, string_label_tokens, nodes, edges, cnt_edge, adj_matrix

    def relation_to_relation_fill(self, node_dict, rel_dict, adj_matrix):
        adj_matrix_temp = np.array(adj_matrix)
        rel_idx_list = []
        for rel in rel_dict.keys():
            rel_idx = node_dict[rel][1]
            rel_idx_list.append(rel_idx)
        adj_matrix_np = np.array(adj_matrix)
        adj_matrix_np_bool = (adj_matrix_np==-1)
        #reassign -1s to 0s
        adj_matrix_np[adj_matrix_np_bool] = 0
        #get squared matrix for r-r
        adj_matrix_sq = adj_matrix_np@adj_matrix_np

        #old adj_matrix + squared matrix only r-r
        rel_idx_list = np.array(rel_idx_list, dtype=np.intp)
        adj_matrix_temp[rel_idx_list[:,np.newaxis], rel_idx_list] = (adj_matrix_sq[rel_idx_list][:,rel_idx_list] > 0)*4
        adj_matrix_new = adj_matrix_temp.tolist()

        return adj_matrix_new

    def get_all_entities_per_sample(self, mark_entity_number, mark_entity, entry):
        text_entity = set()
        text_relation = set()
        for entity_id in mark_entity_number:
            entity = entry['kbs'][entity_id]
            if len(entity[0]) == 0:
                continue
            for rel in entity[2]:
                if len(rel[0]) != 0 and len(rel[1]) != 0:
                    text_relation.add(rel[0])
                    text_entity.add(rel[1])

        text_entity_list = list(text_entity)+list(text_relation)
        text_relation_list = list(text_relation)
        for entity_ele in mark_entity:
            if entity_ele in text_entity_list:
                text_entity_list.remove(entity_ele)

        return text_entity_list, text_relation_list

    def get_change_per_sample(self, mark_entity, text_entity, text_relation):
        # during fine-tuning, we don't mask entities or relations
        ent_change = {}
        total_entity = mark_entity + text_entity

        for ent_id in range(len(total_entity)):
            entity_toks = self.tokenizer.encode(" {}".format(total_entity[ent_id]), add_special_tokens=False)
            ent_change[total_entity[ent_id]] = [entity_toks, ent_id]
        # relation change only includes the relation tokens and ids
        rel_change = {}
        for rel_id in range(len(text_relation)):
            rel_change[text_relation[rel_id]] = self.tokenizer.encode(' {}'.format(text_relation[rel_id]),
                                                                      add_special_tokens=False)
        return ent_change, rel_change

    def truncate_pair_ar(self, a, add_bos_id, graph_ids, text_ids, node_ids, edge_ids):
        # add_bos_id + graph_ids + a + text_ids + b + eos_token_id
        length_a_b = self.args.max_input_length - len(add_bos_id) - len(graph_ids) - len(text_ids) - 1
        if len(a) > length_a_b:
            a = a[:length_a_b]
            node_ids = node_ids[:length_a_b]
            edge_ids = edge_ids[:length_a_b]
        input_ids = add_bos_id + graph_ids + a + text_ids + [self.tokenizer.eos_token_id]
        input_node_ids = [-1] * (len(add_bos_id) + len(graph_ids)) + node_ids + [-1] * (len(text_ids) + 1)
        input_edge_ids = [-1] * (len(add_bos_id) + len(graph_ids)) + edge_ids + [-1] * (len(text_ids) + 1)
        attn_mask = [1] * len(input_ids) + [0] * (self.args.max_input_length - len(input_ids))
        input_ids += [self.tokenizer.pad_token_id] * (self.args.max_input_length - len(input_ids))
        input_node_ids += [-1] * (self.args.max_input_length - len(input_node_ids))
        input_edge_ids += [-1] * (self.args.max_input_length - len(input_edge_ids))
        assert len(input_ids) == len(attn_mask) == self.args.max_input_length == len(input_node_ids) == len(
            input_edge_ids)
        return input_ids, attn_mask, input_node_ids, input_edge_ids

    def ar_prep_data(self, questions, add_bos_id, graph_ids, text_ids, node_ids, edge_ids):
        input_ids, input_attn_mask, input_node_ids, input_edge_ids = self.truncate_pair_ar(questions, add_bos_id,
                                                                                graph_ids, text_ids,
                                                                                node_ids, edge_ids)

        return input_ids, input_attn_mask, input_node_ids, input_edge_ids

    def __getitem__(self, idx):

        entry = self.data[idx]

        entities = []
        for _ in entry['kbs']:
            entities.append(_)

        strings_label = []
        node_ids = []
        edge_ids = []
        strings_label_tokens = ''

        # mark_entity: entities with KB numbers which are important for this task
        # text_entity: entities without KB numbers but only with text, which are less important
        mark_entity = [entry['kbs'][ele_entity][0] for ele_entity in entities]
        mark_entity_number = entities
        text_entity, text_relation = self.get_all_entities_per_sample(mark_entity_number, mark_entity, entry)
        entity_change, relation_change = self.get_change_per_sample(mark_entity, text_entity, text_relation)
        total_entity = mark_entity + text_entity
        adj_matrix = [[-1] * (self.args.max_node_length + 1) for _ in range(self.args.max_node_length + 1)]

        cnt_edge = 0

        if 'title' in entry:
            entity = self.knowledge[entry['title_kb_id']]
            string_label, string_label_tokens, nodes, edges, cnt_edge, adj_matrix = self.linearize_v2(
                entity,
                entity_change,
                self.head_ids,
                self.rel_ids, self.tail_ids,
                relation_change, cnt_edge, adj_matrix)

            strings_label += string_label
            strings_label_tokens += string_label_tokens

        for i, entity_id in enumerate(entities):
            entity = entry['kbs'][entity_id]
            string_label, string_label_tokens, nodes, edges, cnt_edge, adj_matrix = self.linearize_v2(
                entity,
                entity_change,
                self.head_ids,
                self.rel_ids, self.tail_ids,
                relation_change, cnt_edge, adj_matrix)

            strings_label += string_label
            strings_label_tokens += string_label_tokens
            node_ids += nodes
            edge_ids += edges

        if self.topology['relation-relation']:
            adj_matrix = self.relation_to_relation_fill(entity_change, relation_change, adj_matrix)

        words_label_ids, words_label_tokens, words_input_ids, words_input_tokens = [], '', [], ''

        input_ids_ar, attn_mask_ar, input_node_ids_ar, input_edge_ids_ar = \
            self.ar_prep_data(strings_label, self.add_bos_id, self.graph_ids,
                              self.text_ids, node_ids, edge_ids)

        node_length_ar = max(input_node_ids_ar) + 1
        edge_length_ar = max(input_edge_ids_ar) + 1

        def masked_fill(src, masked_value, fill_value):
            return [src[src_id] if src[src_id] != masked_value and src[src_id] < fill_value else fill_value for src_id
                    in range(len(src))]

        input_node_ids_ar, input_edge_ids_ar = masked_fill(input_node_ids_ar, -1, self.args.max_node_length), \
                                               masked_fill(input_edge_ids_ar, -1, self.args.max_edge_length)

        def masked_fill_matrix(adj_matrix_input, masked_value, fill_value):
            adj_matrix_tmp = copy.deepcopy(adj_matrix_input)
            for a_id in range(len(adj_matrix_tmp)):
                for b_id in range(len(adj_matrix_tmp)):
                    if adj_matrix_tmp[a_id][b_id] == masked_value or adj_matrix_tmp[a_id][b_id] > fill_value:
                        adj_matrix_tmp[a_id][b_id] = fill_value
            return adj_matrix_tmp

        adj_matrix_ar = masked_fill_matrix(adj_matrix, -1, self.args.max_edge_length)

        assert len(input_ids_ar) == len(attn_mask_ar) == self.args.max_input_length == len(input_node_ids_ar) == len(
            input_edge_ids_ar)

        input_ids_ar = torch.LongTensor(input_ids_ar)
        attn_mask_ar = torch.LongTensor(attn_mask_ar)

        input_node_ids_ar = torch.LongTensor(input_node_ids_ar)
        input_edge_ids_ar = torch.LongTensor(input_edge_ids_ar)
        node_length_ar = torch.LongTensor([node_length_ar])
        edge_length_ar = torch.LongTensor([edge_length_ar])
        adj_matrix_ar = torch.LongTensor(adj_matrix_ar)

        return input_ids_ar, attn_mask_ar, input_node_ids_ar, node_length_ar, adj_matrix_ar

scoring/aggregate_scorer.py
class ScorerWrapper:
    def __init__(self, scorers, scoring_method="logsum", batch_size=1):
        assert scoring_method in ["product", "logsum"], "Unrecognized `scoring_method`"

        self.scorers = scorers
        self.scoring_method = scoring_method

        # if self.scoring_method == "logsum":
        #     self.score_func = logsum_score
        # elif self.scoring_method == "product":
        #     self.score_func = product_score

        if batch_size > 1:
            exec("self.score_func = {}".format(self.scoring_method+"_"+"score_batched"))
        else:
            exec("self.score_func = {}").format(self.scoring_method+"_"+"score")
        self.batch_size = batch_size
    def get_score_names(self):
        return [s["name"] for s in self.scorers]

    def score_batched(self, input_texts=None, generated_texts=None, old_kgs=None, new_kgs=None, dels_ents=None, partial=False, printing=False, timings=False, extras={}, progress=False):
        assert len(input_texts) == len(generated_texts) == len(old_kgs) == len(new_kgs) == len(dels_ents), "Data lengths don't match"

        data_list = []
        for inp, gen, old_kg, new_kg, del_ents in zip(input_texts, generated_texts, old_kgs, new_kgs, dels_ents):
            data_list.append({"inp": inp, "gen": gen, "old_kg": old_kg, "new_kg": new_kg, "del_ents": del_ents})

        if len(data_list) == 0:
            progress = False

        for batch in batcher(data_list, batch_size=self.batch_size, progress=progress):
            batch_inputs = [instance_dict["inp"] for instance_dict in batch]
            batch_gens = [instance_dict["gen"] for instance_dict in batch]
            batch_old_kgs = [instance_dict["old_kg"] for instance_dict in batch]
            batch_new_kgs = [instance_dict["new_kg"] for instance_dict in batch]
            batch_dels_ents = [instance_dict["del_ents"] for instance_dict in batch]
            batch_scores = self.score_func(self.scorers, batch_inputs, batch_gens, batch_old_kgs, batch_new_kgs, batch_dels_ents)
            for score_type, scores in batch_scores.items():
                if type(scores) in [torch.Tensor, np.array, np.ndarray]:
                    batch_scores[score_type] = scores.tolist()

        if printing:
            print("[total]", all_outputs["total_scores"])
        return batch_scores

    def score(self, input_text=None, generated_text=None, old_kg=None, new_kg=None, del_ents=None):
        aggregate_score = self.score_func(self.scorers, input_text, generated_text, old_kg, new_kg, del_ents)
        return aggregate_score

    def __call__(self, graphs, input_text, generated_text, **kwargs):
        return self.score(graphs, input_text, generated_text, **kwargs)

GAP/modeling_gap.py
class GAPBartForConditionalGeneration(BartForConditionalGeneration):
    def __init__(self, config):
        super().__init__(config)
        base_model = GAPBartModel(config)
        self.model = base_model
        self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))

    def forward(self, input_ids, attention_mask=None, encoder_outputs=None,
            decoder_input_ids=None, decoder_attention_mask=None, input_node_ids=None,
            node_length=None, adj_matrix=None, decoder_whole_ids=None, decoder_cached_states=None,
            use_cache=False, is_training=False):

        if is_training:
            _decoder_input_ids = shift_tokens_right(decoder_input_ids, self.config.pad_token_id)
        else:
            _decoder_input_ids = decoder_input_ids

        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            encoder_outputs=encoder_outputs,
            decoder_input_ids=_decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            input_node_ids=input_node_ids,
            node_length=node_length,
            adj_matrix=adj_matrix,
            decoder_cached_states=decoder_cached_states,
            use_cache=use_cache,
        )
        lm_logits = F.linear(outputs[0], self.model.shared.weight, bias=self.final_logits_bias)
        if is_training:
            loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
            loss = loss_fct(lm_logits.view(-1, self.config.vocab_size),
                              decoder_input_ids.view(-1))
            return loss
        return (lm_logits, ) + outputs[1:]

    @torch.no_grad()
    def generate(
            self,
            input_ids: Optional[torch.LongTensor] = None,
            max_length: Optional[int] = None,
            min_length: Optional[int] = None,
            do_sample: Optional[bool] = None,
            early_stopping: Optional[bool] = None,
            num_beams: Optional[int] = None,
            temperature: Optional[float] = None,
            top_k: Optional[int] = None,
            top_p: Optional[float] = None,
            repetition_penalty: Optional[float] = None,
            bad_words_ids: Optional[Iterable[int]] = None,
            bos_token_id: Optional[int] = None,
            pad_token_id: Optional[int] = None,
            eos_token_id: Optional[int] = None,
            length_penalty: Optional[float] = None,
            no_repeat_ngram_size: Optional[int] = None,
            num_return_sequences: Optional[int] = None,
            attention_mask: Optional[torch.LongTensor] = None,
            input_node_ids=None,
            node_length=None,
            adj_matrix=None,
            decoder_start_token_id: Optional[int] = None,
            use_cache: Optional[bool] = None,
            **model_specific_kwargs
    ) -> torch.LongTensor:
        r""" Generates sequences for models with a LM head. The method currently supports greedy decoding, beam-search decoding, sampling with temperature, sampling with top-k or nucleus sampling.

        Adapted in part from `Facebook's XLM beam search code`_.

        .. _`Facebook's XLM beam search code`:
           https://github.com/facebookresearch/XLM/blob/9e6f6814d17be4fe5b15f2e6c43eb2b2d76daeb4/src/model/transformer.py#L529

        Parameters:

            input_ids: (`optional`) `torch.LongTensor` of shape `(batch_size, sequence_length)`
                The sequence used as a prompt for the generation. If `None` the method initializes
                it as an empty `torch.LongTensor` of shape `(1,)`.

            max_length: (`optional`) int
                The max length of the sequence to be generated.  Between `min_length` and infinity. Default to 20.

            min_length: (`optional`) int
                The min length of the sequence to be generated.  Between 0 and infinity. Default to 0.

            do_sample: (`optional`) bool
                If set to `False` greedy decoding is used. Otherwise sampling is used. Defaults to `False` as defined in `configuration_utils.PretrainedConfig`.

            early_stopping: (`optional`) bool
                if set to `True` beam search is stopped when at least `num_beams` sentences finished per batch. Defaults to `False` as defined in `configuration_utils.PretrainedConfig`.

            num_beams: (`optional`) int
                Number of beams for beam search. Must be between 1 and infinity. 1 means no beam search. Default to 1.

            temperature: (`optional`) float
                The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.

            top_k: (`optional`) int
                The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.

            top_p: (`optional`) float
                The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.

            repetition_penalty: (`optional`) float
                The parameter for repetition penalty. Between 1.0 and infinity. 1.0 means no penalty. Default to 1.0.

            pad_token_id: (`optional`) int
                Padding token. Default to specicic model pad_token_id or None if it does not exist.

            bos_token_id: (`optional`) int
                BOS token. Defaults to `bos_token_id` as defined in the models config.

            eos_token_id: (`optional`) int
                EOS token. Defaults to `eos_token_id` as defined in the models config.

            length_penalty: (`optional`) float
                Exponential penalty to the length. Default to 1.

            no_repeat_ngram_size: (`optional`) int
                If set to int > 0, all ngrams of size `no_repeat_ngram_size` can only occur once.
            bad_words_ids: (`optional`) list of lists of int
                `bad_words_ids` contains tokens that are not allowed to be generated. In order to get the tokens of the words that should not appear in the generated text, use `tokenizer.encode(bad_word, add_prefix_space=True)`.

            num_return_sequences: (`optional`) int
                The number of independently computed returned sequences for each element in the batch. Default to 1.

            attention_mask (`optional`) obj: `torch.LongTensor` of same shape as `input_ids`
                Mask to avoid performing attention on padding token indices.
                Mask values selected in ``[0, 1]``:
                ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
                Defaults to `None`.

                `What are attention masks? <../glossary.html#attention-mask>`__

            decoder_start_token_id=None: (`optional`) int
                Start token id for the decoder. Defaults to ``decoder_start_token_id`` as defined the model's config or to the ``bos_token_id``
                if no ``decoder_start_token_id`` is found in the config.
                This is only relevant for encoder-decoder models.

            use_cache: (`optional`) bool
                If `use_cache` is True, past key values are used to speed up decoding if applicable to model. Defaults to `True`.

            model_specific_kwargs: (`optional`) dict
                Additional model specific kwargs will be forwarded to the `forward` function of the model.

        Return:

            output: `torch.LongTensor` of shape `(batch_size * num_return_sequences, sequence_length)`
                sequence_length is either equal to max_length or shorter if all batches finished early due to the `eos_token_id`

        Examples::

            from transformers import AutoTokenizer, AutoModelForCausalLM

            tokenizer = AutoTokenizer.from_pretrained('distilgpt2')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('distilgpt2')    # Download model and configuration from S3 and cache.
            outputs = model.generate(max_length=40)  # do greedy decoding
            print('Generated: {}'.format(tokenizer.decode(outputs[0], skip_special_tokens=True)))

            tokenizer = AutoTokenizer.from_pretrained('openai-gpt')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('openai-gpt')    # Download model and configuration from S3 and cache.
            input_context = 'The dog'
            input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
            outputs = model.generate(input_ids=input_ids, num_beams=5, num_return_sequences=3, temperature=1.5)  # generate 3 independent sequences using beam search decoding (5 beams) with sampling from initial context 'The dog'
            for i in range(3): #  3 output sequences were generated
                print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))

            tokenizer = AutoTokenizer.from_pretrained('distilgpt2')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('distilgpt2')    # Download model and configuration from S3 and cache.
            input_context = 'The dog'
            input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
            outputs = model.generate(input_ids=input_ids, max_length=40, temperature=0.7, num_return_sequences=3, do_sample=True)  # 3 generate sequences using by sampling
            for i in range(3): #  3 output sequences were generated
                print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))

            tokenizer = AutoTokenizer.from_pretrained('ctrl')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('ctrl')    # Download model and configuration from S3 and cache.
            input_context = 'Legal My neighbor is'  # "Legal" is one of the control codes for ctrl
            input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
            outputs = model.generate(input_ids=input_ids, max_length=50, temperature=0.7, repetition_penalty=1.2)  # generate sequences
            print('Generated: {}'.format(tokenizer.decode(outputs[0], skip_special_tokens=True)))

            tokenizer = AutoTokenizer.from_pretrained('gpt2')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('gpt2')    # Download model and configuration from S3 and cache.
            input_context = 'My cute dog'  # "Legal" is one of the control codes for ctrl
            bad_words_ids = [tokenizer.encode(bad_word, add_prefix_space=True) for bad_word in ['idiot', 'stupid', 'shut up']]
            input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
            outputs = model.generate(input_ids=input_ids, max_length=100, do_sample=True, bad_words_ids=bad_words_ids)  # generate sequences without allowing bad_words to be generated
        """

        # We cannot generate if the model does not have a LM head
        if self.get_output_embeddings() is None:
            raise AttributeError(
                "You tried to generate sequences with a model that does not have a LM Head."
                "Please use another model class (e.g. `OpenAIGPTLMHeadModel`, `XLNetLMHeadModel`, `GPT2LMHeadModel`, `CTRLLMHeadModel`, `T5WithLMHeadModel`, `TransfoXLLMHeadModel`, `XLMWithLMHeadModel`, `BartForConditionalGeneration` )"
            )

        max_length = max_length if max_length is not None else self.config.max_length
        min_length = min_length if min_length is not None else self.config.min_length
        do_sample = do_sample if do_sample is not None else self.config.do_sample
        early_stopping = early_stopping if early_stopping is not None else self.config.early_stopping
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        num_beams = num_beams if num_beams is not None else self.config.num_beams
        temperature = temperature if temperature is not None else self.config.temperature
        top_k = top_k if top_k is not None else self.config.top_k
        top_p = top_p if top_p is not None else self.config.top_p
        repetition_penalty = repetition_penalty if repetition_penalty is not None else self.config.repetition_penalty
        bos_token_id = bos_token_id if bos_token_id is not None else self.config.bos_token_id
        pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id
        eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
        length_penalty = length_penalty if length_penalty is not None else self.config.length_penalty
        no_repeat_ngram_size = (
            no_repeat_ngram_size if no_repeat_ngram_size is not None else self.config.no_repeat_ngram_size
        )
        bad_words_ids = bad_words_ids if bad_words_ids is not None else self.config.bad_words_ids
        num_return_sequences = (
            num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
        )
        decoder_start_token_id = (
            decoder_start_token_id if decoder_start_token_id is not None else self.config.decoder_start_token_id
        )

        if input_ids is not None:
            batch_size = input_ids.shape[0]  # overriden by the input batch_size
        else:
            batch_size = 1

        assert isinstance(max_length, int) and max_length > 0, "`max_length` should be a strictly positive integer."
        assert isinstance(min_length, int) and min_length >= 0, "`min_length` should be a positive integer."
        assert isinstance(do_sample, bool), "`do_sample` should be a boolean."
        assert isinstance(early_stopping, bool), "`early_stopping` should be a boolean."
        assert isinstance(use_cache, bool), "`use_cache` should be a boolean."
        assert isinstance(num_beams, int) and num_beams > 0, "`num_beams` should be a strictly positive integer."
        assert temperature > 0, "`temperature` should be strictly positive."
        assert isinstance(top_k, int) and top_k >= 0, "`top_k` should be a positive integer."
        assert 0 <= top_p <= 1, "`top_p` should be between 0 and 1."
        assert repetition_penalty >= 1.0, "`repetition_penalty` should be >= 1."
        assert input_ids is not None or (
                isinstance(bos_token_id, int) and bos_token_id >= 0
        ), "If input_ids is not defined, `bos_token_id` should be a positive integer."
        assert pad_token_id is None or (
                isinstance(pad_token_id, int) and (pad_token_id >= 0)
        ), "`pad_token_id` should be a positive integer."
        assert (eos_token_id is None) or (
                isinstance(eos_token_id, int) and (eos_token_id >= 0)
        ), "`eos_token_id` should be a positive integer."
        assert length_penalty > 0, "`length_penalty` should be strictly positive."
        assert (
                isinstance(no_repeat_ngram_size, int) and no_repeat_ngram_size >= 0
        ), "`no_repeat_ngram_size` should be a positive integer."
        assert (
                isinstance(num_return_sequences, int) and num_return_sequences > 0
        ), "`num_return_sequences` should be a strictly positive integer."
        assert (
                bad_words_ids is None or isinstance(bad_words_ids, list) and isinstance(bad_words_ids[0], list)
        ), "`bad_words_ids` is either `None` or a list of lists of tokens that should not be generated"

        if input_ids is None:
            assert isinstance(bos_token_id, int) and bos_token_id >= 0, (
                "you should either supply a context to complete as `input_ids` input "
                "or a `bos_token_id` (integer >= 0) as a first token to start the generation."
            )
            input_ids = torch.full(
                (batch_size, 1), bos_token_id, dtype=torch.long, device=next(self.parameters()).device,
            )
        else:
            assert input_ids.dim() == 2, "Input prompt should be of shape (batch_size, sequence length)."

        # not allow to duplicate outputs when greedy decoding
        if do_sample is False:
            if num_beams == 1:
                # no_beam_search greedy generation conditions
                assert (
                        num_return_sequences == 1
                ), "Greedy decoding will always produce the same output for num_beams == 1 and num_return_sequences > 1. Please set num_return_sequences = 1"

            else:
                # beam_search greedy generation conditions
                assert (
                        num_beams >= num_return_sequences
                ), "Greedy beam search decoding cannot return more sequences than it has beams. Please set num_beams >= num_return_sequences"

        # create attention mask if necessary
        # TODO (PVP): this should later be handled by the forward fn() in each model in the future see PR 3140
        if (attention_mask is None) and (pad_token_id is not None) and (pad_token_id in input_ids):
            attention_mask = input_ids.ne(pad_token_id).long()
        elif attention_mask is None:
            attention_mask = input_ids.new_ones(input_ids.shape)

        # set pad_token_id to eos_token_id if not set. Important that this is done after
        # attention_mask is created
        if pad_token_id is None and eos_token_id is not None:
            logger.warning(
                "Setting `pad_token_id` to {} (first `eos_token_id`) to generate sequence".format(eos_token_id)
            )
            pad_token_id = eos_token_id

        # current position and vocab size
        if hasattr(self.config, "vocab_size"):
            vocab_size = self.config.vocab_size
        elif (
                self.config.is_encoder_decoder
                and hasattr(self.config, "decoder")
                and hasattr(self.config.decoder, "vocab_size")
        ):
            vocab_size = self.config.decoder.vocab_size

        # set effective batch size and effective batch multiplier according to do_sample
        if do_sample:
            effective_batch_size = batch_size * num_return_sequences
            effective_batch_mult = num_return_sequences
        else:
            effective_batch_size = batch_size
            effective_batch_mult = 1

        if self.config.is_encoder_decoder:
            if decoder_start_token_id is None:
                decoder_start_token_id = bos_token_id

            assert (
                    decoder_start_token_id is not None
            ), "decoder_start_token_id or bos_token_id has to be defined for encoder-decoder generation"
            assert hasattr(self, "get_encoder"), "{} should have a 'get_encoder' function defined".format(self)
            assert callable(self.get_encoder), "{} should be a method".format(self.get_encoder)

            # get encoder and store encoder outputs
            encoder = self.get_encoder()

            # add structural information when encoding
            encoder_outputs: tuple = encoder(input_ids, attention_mask=attention_mask, input_node_ids=input_node_ids,
                                             node_length=node_length, adj_matrix=adj_matrix)

        # Expand input ids if num_beams > 1 or num_return_sequences > 1
        if num_return_sequences > 1 or num_beams > 1:
            input_ids_len = input_ids.shape[-1]
            input_ids = input_ids.unsqueeze(1).expand(batch_size, effective_batch_mult * num_beams, input_ids_len)
            attention_mask = attention_mask.unsqueeze(1).expand(
                batch_size, effective_batch_mult * num_beams, input_ids_len
            )

            input_ids = input_ids.contiguous().view(
                effective_batch_size * num_beams, input_ids_len
            )  # shape: (batch_size * num_return_sequences * num_beams, cur_len)
            attention_mask = attention_mask.contiguous().view(
                effective_batch_size * num_beams, input_ids_len
            )  # shape: (batch_size * num_return_sequences * num_beams, cur_len)

        if self.config.is_encoder_decoder:
            # create empty decoder_input_ids
            input_ids = torch.full(
                (effective_batch_size * num_beams, 1),
                decoder_start_token_id,
                dtype=torch.long,
                device=next(self.parameters()).device,
            )
            cur_len = 1

            assert (
                    batch_size == encoder_outputs[0].shape[0]
            ), f"expected encoder_outputs[0] to have 1st dimension bs={batch_size}, got {encoder_outputs[0].shape[0]} "

            # expand batch_idx to assign correct encoder output for expanded input_ids (due to num_beams > 1 and num_return_sequences > 1)
            expanded_batch_idxs = (
                torch.arange(batch_size)
                    .view(-1, 1)
                    .repeat(1, num_beams * effective_batch_mult)
                    .view(-1)
                    .to(input_ids.device)
            )
            # expand encoder_outputs
            encoder_outputs = (encoder_outputs[0].index_select(0, expanded_batch_idxs), *encoder_outputs[1:])

        else:
            encoder_outputs = None
            cur_len = input_ids.shape[-1]

        if num_beams > 1:
            output = self._generate_beam_search(
                input_ids,
                cur_len=cur_len,
                max_length=max_length,
                min_length=min_length,
                do_sample=do_sample,
                early_stopping=early_stopping,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                bad_words_ids=bad_words_ids,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
                batch_size=effective_batch_size,
                num_return_sequences=num_return_sequences,
                length_penalty=length_penalty,
                num_beams=num_beams,
                vocab_size=vocab_size,
                encoder_outputs=encoder_outputs,
                attention_mask=attention_mask,
                use_cache=use_cache,
                model_specific_kwargs=model_specific_kwargs,
            )
        else:
            output = self._generate_no_beam_search(
                input_ids,
                cur_len=cur_len,
                max_length=max_length,
                min_length=min_length,
                do_sample=do_sample,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                bad_words_ids=bad_words_ids,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
                batch_size=effective_batch_size,
                encoder_outputs=encoder_outputs,
                attention_mask=attention_mask,
                use_cache=use_cache,
                model_specific_kwargs=model_specific_kwargs,
            )

        return output

GAP/modeling_gap_type.py
class GAPBartForConditionalGeneration(BartForConditionalGeneration):
    def __init__(self, config, **kwargs):
        super().__init__(config)
        base_model = GAPBartModel(config,**kwargs)
        self.model = base_model
        self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))

    def forward(self, input_ids, attention_mask=None, encoder_outputs=None,
            decoder_input_ids=None, decoder_attention_mask=None, input_node_ids=None,
            node_length=None, adj_matrix=None, decoder_whole_ids=None, decoder_cached_states=None,
            use_cache=False, is_training=False):

        if is_training:
            _decoder_input_ids = shift_tokens_right(decoder_input_ids, self.config.pad_token_id)
        else:
            _decoder_input_ids = decoder_input_ids

        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            encoder_outputs=encoder_outputs,
            decoder_input_ids=_decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            input_node_ids=input_node_ids,
            node_length=node_length,
            adj_matrix=adj_matrix,
            decoder_cached_states=decoder_cached_states,
            use_cache=use_cache,
        )
        lm_logits = F.linear(outputs[0], self.model.shared.weight, bias=self.final_logits_bias)
        if is_training:
            loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
            loss = loss_fct(lm_logits.view(-1, self.config.vocab_size),
                              decoder_input_ids.view(-1))
            return loss
        return (lm_logits, ) + outputs[1:]

    @torch.no_grad()
    def generate(
            self,
            input_ids: Optional[torch.LongTensor] = None,
            max_length: Optional[int] = None,
            min_length: Optional[int] = None,
            do_sample: Optional[bool] = None,
            early_stopping: Optional[bool] = None,
            num_beams: Optional[int] = None,
            temperature: Optional[float] = None,
            top_k: Optional[int] = None,
            top_p: Optional[float] = None,
            repetition_penalty: Optional[float] = None,
            bad_words_ids: Optional[Iterable[int]] = None,
            bos_token_id: Optional[int] = None,
            pad_token_id: Optional[int] = None,
            eos_token_id: Optional[int] = None,
            length_penalty: Optional[float] = None,
            no_repeat_ngram_size: Optional[int] = None,
            num_return_sequences: Optional[int] = None,
            attention_mask: Optional[torch.LongTensor] = None,
            input_node_ids=None,
            node_length=None,
            adj_matrix=None,
            decoder_start_token_id: Optional[int] = None,
            use_cache: Optional[bool] = None,
            **model_specific_kwargs
    ) -> torch.LongTensor:
        r""" Generates sequences for models with a LM head. The method currently supports greedy decoding, beam-search decoding, sampling with temperature, sampling with top-k or nucleus sampling.

        Adapted in part from `Facebook's XLM beam search code`_.

        .. _`Facebook's XLM beam search code`:
           https://github.com/facebookresearch/XLM/blob/9e6f6814d17be4fe5b15f2e6c43eb2b2d76daeb4/src/model/transformer.py#L529

        Parameters:

            input_ids: (`optional`) `torch.LongTensor` of shape `(batch_size, sequence_length)`
                The sequence used as a prompt for the generation. If `None` the method initializes
                it as an empty `torch.LongTensor` of shape `(1,)`.

            max_length: (`optional`) int
                The max length of the sequence to be generated.  Between `min_length` and infinity. Default to 20.

            min_length: (`optional`) int
                The min length of the sequence to be generated.  Between 0 and infinity. Default to 0.

            do_sample: (`optional`) bool
                If set to `False` greedy decoding is used. Otherwise sampling is used. Defaults to `False` as defined in `configuration_utils.PretrainedConfig`.

            early_stopping: (`optional`) bool
                if set to `True` beam search is stopped when at least `num_beams` sentences finished per batch. Defaults to `False` as defined in `configuration_utils.PretrainedConfig`.

            num_beams: (`optional`) int
                Number of beams for beam search. Must be between 1 and infinity. 1 means no beam search. Default to 1.

            temperature: (`optional`) float
                The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.

            top_k: (`optional`) int
                The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.

            top_p: (`optional`) float
                The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.

            repetition_penalty: (`optional`) float
                The parameter for repetition penalty. Between 1.0 and infinity. 1.0 means no penalty. Default to 1.0.

            pad_token_id: (`optional`) int
                Padding token. Default to specicic model pad_token_id or None if it does not exist.

            bos_token_id: (`optional`) int
                BOS token. Defaults to `bos_token_id` as defined in the models config.

            eos_token_id: (`optional`) int
                EOS token. Defaults to `eos_token_id` as defined in the models config.

            length_penalty: (`optional`) float
                Exponential penalty to the length. Default to 1.

            no_repeat_ngram_size: (`optional`) int
                If set to int > 0, all ngrams of size `no_repeat_ngram_size` can only occur once.
            bad_words_ids: (`optional`) list of lists of int
                `bad_words_ids` contains tokens that are not allowed to be generated. In order to get the tokens of the words that should not appear in the generated text, use `tokenizer.encode(bad_word, add_prefix_space=True)`.

            num_return_sequences: (`optional`) int
                The number of independently computed returned sequences for each element in the batch. Default to 1.

            attention_mask (`optional`) obj: `torch.LongTensor` of same shape as `input_ids`
                Mask to avoid performing attention on padding token indices.
                Mask values selected in ``[0, 1]``:
                ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
                Defaults to `None`.

                `What are attention masks? <../glossary.html#attention-mask>`__

            decoder_start_token_id=None: (`optional`) int
                Start token id for the decoder. Defaults to ``decoder_start_token_id`` as defined the model's config or to the ``bos_token_id``
                if no ``decoder_start_token_id`` is found in the config.
                This is only relevant for encoder-decoder models.

            use_cache: (`optional`) bool
                If `use_cache` is True, past key values are used to speed up decoding if applicable to model. Defaults to `True`.

            model_specific_kwargs: (`optional`) dict
                Additional model specific kwargs will be forwarded to the `forward` function of the model.

        Return:

            output: `torch.LongTensor` of shape `(batch_size * num_return_sequences, sequence_length)`
                sequence_length is either equal to max_length or shorter if all batches finished early due to the `eos_token_id`

        Examples::

            from transformers import AutoTokenizer, AutoModelForCausalLM

            tokenizer = AutoTokenizer. ('distilgpt2')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('distilgpt2')    # Download model and configuration from S3 and cache.
            outputs = model.generate(max_length=40)  # do greedy decoding
            print('Generated: {}'.format(tokenizer.decode(outputs[0], skip_special_tokens=True)))

            tokenizer = AutoTokenizer.from_pretrained('openai-gpt')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('openai-gpt')    # Download model and configuration from S3 and cache.
            input_context = 'The dog'
            input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
            outputs = model.generate(input_ids=input_ids, num_beams=5, num_return_sequences=3, temperature=1.5)  # generate 3 independent sequences using beam search decoding (5 beams) with sampling from initial context 'The dog'
            for i in range(3): #  3 output sequences were generated
                print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))

            tokenizer = AutoTokenizer.from_pretrained('distilgpt2')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('distilgpt2')    # Download model and configuration from S3 and cache.
            input_context = 'The dog'
            input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
            outputs = model.generate(input_ids=input_ids, max_length=40, temperature=0.7, num_return_sequences=3, do_sample=True)  # 3 generate sequences using by sampling
            for i in range(3): #  3 output sequences were generated
                print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))

            tokenizer = AutoTokenizer.from_pretrained('ctrl')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('ctrl')    # Download model and configuration from S3 and cache.
            input_context = 'Legal My neighbor is'  # "Legal" is one of the control codes for ctrl
            input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
            outputs = model.generate(input_ids=input_ids, max_length=50, temperature=0.7, repetition_penalty=1.2)  # generate sequences
            print('Generated: {}'.format(tokenizer.decode(outputs[0], skip_special_tokens=True)))

            tokenizer = AutoTokenizer.from_pretrained('gpt2')   # Initialize tokenizer
            model = AutoModelForCausalLM.from_pretrained('gpt2')    # Download model and configuration from S3 and cache.
            input_context = 'My cute dog'  # "Legal" is one of the control codes for ctrl
            bad_words_ids = [tokenizer.encode(bad_word, add_prefix_space=True) for bad_word in ['idiot', 'stupid', 'shut up']]
            input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
            outputs = model.generate(input_ids=input_ids, max_length=100, do_sample=True, bad_words_ids=bad_words_ids)  # generate sequences without allowing bad_words to be generated
        """

        # We cannot generate if the model does not have a LM head
        if self.get_output_embeddings() is None:
            raise AttributeError(
                "You tried to generate sequences with a model that does not have a LM Head."
                "Please use another model class (e.g. `OpenAIGPTLMHeadModel`, `XLNetLMHeadModel`, `GPT2LMHeadModel`, `CTRLLMHeadModel`, `T5WithLMHeadModel`, `TransfoXLLMHeadModel`, `XLMWithLMHeadModel`, `BartForConditionalGeneration` )"
            )

        max_length = max_length if max_length is not None else self.config.max_length
        min_length = min_length if min_length is not None else self.config.min_length
        do_sample = do_sample if do_sample is not None else self.config.do_sample
        early_stopping = early_stopping if early_stopping is not None else self.config.early_stopping
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        num_beams = num_beams if num_beams is not None else self.config.num_beams
        temperature = temperature if temperature is not None else self.config.temperature
        top_k = top_k if top_k is not None else self.config.top_k
        top_p = top_p if top_p is not None else self.config.top_p
        repetition_penalty = repetition_penalty if repetition_penalty is not None else self.config.repetition_penalty
        bos_token_id = bos_token_id if bos_token_id is not None else self.config.bos_token_id
        pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id
        eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
        length_penalty = length_penalty if length_penalty is not None else self.config.length_penalty
        no_repeat_ngram_size = (
            no_repeat_ngram_size if no_repeat_ngram_size is not None else self.config.no_repeat_ngram_size
        )
        bad_words_ids = bad_words_ids if bad_words_ids is not None else self.config.bad_words_ids
        num_return_sequences = (
            num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
        )
        decoder_start_token_id = (
            decoder_start_token_id if decoder_start_token_id is not None else self.config.decoder_start_token_id
        )

        if input_ids is not None:
            batch_size = input_ids.shape[0]  # overriden by the input batch_size
        else:
            batch_size = 1

        assert isinstance(max_length, int) and max_length > 0, "`max_length` should be a strictly positive integer."
        assert isinstance(min_length, int) and min_length >= 0, "`min_length` should be a positive integer."
        assert isinstance(do_sample, bool), "`do_sample` should be a boolean."
        assert isinstance(early_stopping, bool), "`early_stopping` should be a boolean."
        assert isinstance(use_cache, bool), "`use_cache` should be a boolean."
        assert isinstance(num_beams, int) and num_beams > 0, "`num_beams` should be a strictly positive integer."
        assert temperature > 0, "`temperature` should be strictly positive."
        assert isinstance(top_k, int) and top_k >= 0, "`top_k` should be a positive integer."
        assert 0 <= top_p <= 1, "`top_p` should be between 0 and 1."
        assert repetition_penalty >= 1.0, "`repetition_penalty` should be >= 1."
        assert input_ids is not None or (
                isinstance(bos_token_id, int) and bos_token_id >= 0
        ), "If input_ids is not defined, `bos_token_id` should be a positive integer."
        assert pad_token_id is None or (
                isinstance(pad_token_id, int) and (pad_token_id >= 0)
        ), "`pad_token_id` should be a positive integer."
        assert (eos_token_id is None) or (
                isinstance(eos_token_id, int) and (eos_token_id >= 0)
        ), "`eos_token_id` should be a positive integer."
        assert length_penalty > 0, "`length_penalty` should be strictly positive."
        assert (
                isinstance(no_repeat_ngram_size, int) and no_repeat_ngram_size >= 0
        ), "`no_repeat_ngram_size` should be a positive integer."
        assert (
                isinstance(num_return_sequences, int) and num_return_sequences > 0
        ), "`num_return_sequences` should be a strictly positive integer."
        assert (
                bad_words_ids is None or isinstance(bad_words_ids, list) and isinstance(bad_words_ids[0], list)
        ), "`bad_words_ids` is either `None` or a list of lists of tokens that should not be generated"

        if input_ids is None:
            assert isinstance(bos_token_id, int) and bos_token_id >= 0, (
                "you should either supply a context to complete as `input_ids` input "
                "or a `bos_token_id` (integer >= 0) as a first token to start the generation."
            )
            input_ids = torch.full(
                (batch_size, 1), bos_token_id, dtype=torch.long, device=next(self.parameters()).device,
            )
        else:
            assert input_ids.dim() == 2, "Input prompt should be of shape (batch_size, sequence length)."

        # not allow to duplicate outputs when greedy decoding
        if do_sample is False:
            if num_beams == 1:
                # no_beam_search greedy generation conditions
                assert (
                        num_return_sequences == 1
                ), "Greedy decoding will always produce the same output for num_beams == 1 and num_return_sequences > 1. Please set num_return_sequences = 1"

            else:
                # beam_search greedy generation conditions
                assert (
                        num_beams >= num_return_sequences
                ), "Greedy beam search decoding cannot return more sequences than it has beams. Please set num_beams >= num_return_sequences"

        # create attention mask if necessary
        # TODO (PVP): this should later be handled by the forward fn() in each model in the future see PR 3140
        if (attention_mask is None) and (pad_token_id is not None) and (pad_token_id in input_ids):
            attention_mask = input_ids.ne(pad_token_id).long()
        elif attention_mask is None:
            attention_mask = input_ids.new_ones(input_ids.shape)

        # set pad_token_id to eos_token_id if not set. Important that this is done after
        # attention_mask is created
        if pad_token_id is None and eos_token_id is not None:
            logger.warning(
                "Setting `pad_token_id` to {} (first `eos_token_id`) to generate sequence".format(eos_token_id)
            )
            pad_token_id = eos_token_id

        # current position and vocab size
        if hasattr(self.config, "vocab_size"):
            vocab_size = self.config.vocab_size
        elif (
                self.config.is_encoder_decoder
                and hasattr(self.config, "decoder")
                and hasattr(self.config.decoder, "vocab_size")
        ):
            vocab_size = self.config.decoder.vocab_size

        # set effective batch size and effective batch multiplier according to do_sample
        if do_sample:
            effective_batch_size = batch_size * num_return_sequences
            effective_batch_mult = num_return_sequences
        else:
            effective_batch_size = batch_size
            effective_batch_mult = 1

        if self.config.is_encoder_decoder:
            if decoder_start_token_id is None:
                decoder_start_token_id = bos_token_id

            assert (
                    decoder_start_token_id is not None
            ), "decoder_start_token_id or bos_token_id has to be defined for encoder-decoder generation"
            assert hasattr(self, "get_encoder"), "{} should have a 'get_encoder' function defined".format(self)
            assert callable(self.get_encoder), "{} should be a method".format(self.get_encoder)

            # get encoder and store encoder outputs
            encoder = self.get_encoder()

            # add structural information when encoding
            encoder_outputs: tuple = encoder(input_ids, attention_mask=attention_mask, input_node_ids=input_node_ids,
                                             node_length=node_length, adj_matrix=adj_matrix)

        # Expand input ids if num_beams > 1 or num_return_sequences > 1
        if num_return_sequences > 1 or num_beams > 1:
            input_ids_len = input_ids.shape[-1]
            input_ids = input_ids.unsqueeze(1).expand(batch_size, effective_batch_mult * num_beams, input_ids_len)
            attention_mask = attention_mask.unsqueeze(1).expand(
                batch_size, effective_batch_mult * num_beams, input_ids_len
            )

            input_ids = input_ids.contiguous().view(
                effective_batch_size * num_beams, input_ids_len
            )  # shape: (batch_size * num_return_sequences * num_beams, cur_len)
            attention_mask = attention_mask.contiguous().view(
                effective_batch_size * num_beams, input_ids_len
            )  # shape: (batch_size * num_return_sequences * num_beams, cur_len)

        if self.config.is_encoder_decoder:
            # create empty decoder_input_ids
            input_ids = torch.full(
                (effective_batch_size * num_beams, 1),
                decoder_start_token_id,
                dtype=torch.long,
                device=next(self.parameters()).device,
            )
            cur_len = 1

            assert (
                    batch_size == encoder_outputs[0].shape[0]
            ), f"expected encoder_outputs[0] to have 1st dimension bs={batch_size}, got {encoder_outputs[0].shape[0]} "

            # expand batch_idx to assign correct encoder output for expanded input_ids (due to num_beams > 1 and num_return_sequences > 1)
            expanded_batch_idxs = (
                torch.arange(batch_size)
                    .view(-1, 1)
                    .repeat(1, num_beams * effective_batch_mult)
                    .view(-1)
                    .to(input_ids.device)
            )
            # expand encoder_outputs
            encoder_outputs = (encoder_outputs[0].index_select(0, expanded_batch_idxs), *encoder_outputs[1:])

        else:
            encoder_outputs = None
            cur_len = input_ids.shape[-1]

        if num_beams > 1:
            output = self._generate_beam_search(
                input_ids,
                cur_len=cur_len,
                max_length=max_length,
                min_length=min_length,
                do_sample=do_sample,
                early_stopping=early_stopping,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                bad_words_ids=bad_words_ids,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
                batch_size=effective_batch_size,
                num_return_sequences=num_return_sequences,
                length_penalty=length_penalty,
                num_beams=num_beams,
                vocab_size=vocab_size,
                encoder_outputs=encoder_outputs,
                attention_mask=attention_mask,
                use_cache=use_cache,
                model_specific_kwargs=model_specific_kwargs,
            )
        else:
            output = self._generate_no_beam_search(
                input_ids,
                cur_len=cur_len,
                max_length=max_length,
                min_length=min_length,
                do_sample=do_sample,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                bad_words_ids=bad_words_ids,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
                batch_size=effective_batch_size,
                encoder_outputs=encoder_outputs,
                attention_mask=attention_mask,
                use_cache=use_cache,
                model_specific_kwargs=model_specific_kwargs,
            )

        return output

Answer: model = T5.from_pretrained(checkpoint)

<think>
The function run() is loading a model based on args. After handling type_encoding, the else branch checks if the model path contains 't5'. The next step is to instantiate the appropriate model class. Since T5 is imported as T5ForConditionalGeneration alias T5, the next line should load a T5 model from the checkpoint.
</think>
model = T5.from_pretrained(checkpoint)


Input:
from typing import Dict
from typing import Type
from tasks.affect import ActivityAnalysis
from tasks.affect import ActivityGet
from tasks.affect import SleepAnalysis
from tasks.affect import SleepGet
from tasks.ask_user import AskUser
from tasks.google_translator import GoogleTranslate
from tasks.playwright import Click
from tasks.playwright import CurrentWebPage
from tasks.playwright import ExtractHyperlinks
from tasks.playwright import ExtractText
from tasks.playwright import GetElements
from tasks.playwright import Navigate
from tasks.playwright import NavigateBack
from tasks.read_from_datapipe import ReadDataPipe
from tasks.serpapi import SerpAPI
from tasks.task import BaseTask
from tasks.task_types import TaskType
from tasks.test_file import TestFile

Context:
tasks/types.py

TASK_TO_CLASS: Dict[TaskType, Type[BaseTask]] = {
    TaskType.SERPAPI: SerpAPI,
    TaskType.CLICK: Click,
    TaskType.GET_CURRENT_PAGE: CurrentWebPage,
    TaskType.EXTRACT_HYPERLINKS: ExtractHyperlinks,
    TaskType.EXTRACT_TEXT: ExtractText,
    TaskType.GET_ELEMENTS: GetElements,
    TaskType.NAVIGATE_BACK: NavigateBack,
    TaskType.NAVIGATE: Navigate,
    TaskType.AFFECT_SLEEP_GET: SleepGet,
    TaskType.AFFECT_ACTIVITY_GET: ActivityGet,
    TaskType.AFFECT_SLEEP_ANALYSIS: SleepAnalysis,
    TaskType.AFFECT_ACTIVITY_ANALYSIS: ActivityAnalysis,

tasks/playwright/extract_hyperlinks.py
class ExtractHyperlinks(BaseBrowser):
    """
    **Description:**

        This task extracts all hyperlinks from the current webpage.
    """

    name: str = "extract_hyperlinks"
    chat_name: str = "ExtractHyperLinks"
    description: str = "Extract all hyperlinks on the current webpage"
    dependencies: List[str] = []
    inputs: List[str] = [
        "Boolean: True/False. Return absolute URLs instead of relative URLs."
    ]
    outputs: List[str] = []
    output_type: bool = False

    @model_validator(mode="before")
    def check_bs_import(cls, values: dict) -> dict:
        """
            Check that the arguments are valid.

        Args:
            values (Dict): The current attribute values.
        Return:
            Dict: The updated attribute values.
        Raise:
            ImportError: If 'beautifulsoup4' package is not installed.

        """

        try:
            from bs4 import BeautifulSoup  # noqa: F401
        except ImportError:
            raise ImportError(
                "The 'beautifulsoup4' package is required to use this tool."
                " Please install it with 'pip install beautifulsoup4'."
            )
        return values

    @staticmethod
    def scrape_page(
        page: Any, html_content: str, absolute_urls: bool
    ) -> str:
        """
            Scrape hyperlinks from the current webpage.

        Args:
            page (Any): The current webpage.
            html_content (str): The HTML content of the webpage.
            absolute_urls (bool): True if absolute URLs should be returned, False otherwise.
        Return:
            str: JSON string containing the extracted hyperlinks.

        """

        from urllib.parse import urljoin
        from bs4 import BeautifulSoup

        # Parse the HTML content with BeautifulSoup
        soup = BeautifulSoup(html_content, "lxml")

        # Find all the anchor elements and extract their href attributes
        anchors = soup.find_all("a")
        if absolute_urls:
            base_url = page.url
            links = [
                urljoin(base_url, anchor.get("href", ""))
                for anchor in anchors
            ]
        else:
            links = [anchor.get("href", "") for anchor in anchors]
        # Return the list of links as a JSON string
        return json.dumps(links)

    def _execute(
        self,
        inputs: List[Any],
    ) -> str:
        """
            Execute the ExtractHyperlinks task.

        Args:
            input (str): The input parameter for the task.
        Return:
            str: JSON string containing the extracted hyperlinks.
        Raise:
            ValueError: If the synchronous browser is not provided.

        """
        if self.sync_browser is None:
            raise ValueError(
                f"Synchronous browser not provided to {self.name}"
            )
        page = get_current_page(self.sync_browser)
        html_content = page.content()
        return self.scrape_page(page, html_content, inputs[0])

    def explain(
        self,
    ) -> str:
        """
            Provide a brief explanation of the ExtractHyperlinks task.

        Return:
            str: An explanation of the task.

        """

        return "This task extracts all of the hyperlinks."

tasks/affect/sleep_analysis.py
class SleepAnalysis(Affect):
    """
    **Description:**

        This tasks performs average, sum, or trend analysis on the provided raw sleep affect data for specific patient.
    """

    name: str = "affect_sleep_analysis"
    chat_name: str = "AffectSleepAnalysis"
    description: str = (
        "Performs trend or average analysis on the provided sleep data. You must Call this whenever sleep trend or average is needed."
        "For example, if the user asks for trends (or variations) in data, you must call this task"
    )
    dependencies: List[str] = ["affect_sleep_get"]
    inputs: List[str] = [
        "datapipe key to the data",
        "analysis_type. It can be one of [average, trend].",
    ]
    outputs: List[str] = [
        (
            "The analysis result for total_sleep_time. Look for analysis_type to find the type of analysis. "
            "total_sleep_time (in minutes) is Total amount of sleep (a.k.a. sleep duration) registered during the sleep period."
        ),
        (
            "The analysis result for awake_duration. Look for analysis_type to find the type of analysis. "
            "awake_duration (in minutes) is the total amount of awake time registered during the sleep period."
        ),
        (
            "The analysis result for light_sleep_duration. Look for analysis_type to find the type of analysis. "
            "light_sleep_duration (in minutes) is the total amount of light (N1 or N2) sleep registered during the sleep period."
        ),
        (
            "The analysis result for rem_sleep_duration. Look for analysis_type to find the type of analysis. "
            "rem_sleep_duration (in minutes) is the total amount of REM sleep registered during the sleep period."
        ),
        (
            "The analysis result for deep_sleep_duration. Look for analysis_type to find the type of analysis. "
            "deep_sleep_duration (in minutes) is the total amount of deep (N3) sleep registered during the sleep period."
        ),
        (
            "The analysis result for sleep_onset_latency. Look for analysis_type to find the type of analysis. sleep_onset_latency (in minutes) "
            "is the detected latency from bedtime_start to the beginning of the first five minutes of persistent sleep."
        ),
        (
            "The analysis result for midpoint_time_of_sleep. Look for analysis_type to find the type of analysis. "
            "midpoint_time_of_sleep (in minutes) is the time from the start of sleep to the midpoint of sleep. The midpoint ignores awake periods."
        ),
        (
            "The analysis result for sleep_efficiency. Look for analysis_type to find the type of analysis. "
            "sleep_efficiency is the percentage of the sleep period spent asleep (100% * sleep duration / time in bed)."
        ),
        (
            "The analysis result for average_heart_rate. Look for analysis_type to find the type of analysis. "
            "average_heart_rate is the average heart rate registered during the sleep period."
        ),
        (
            "The analysis result for minimum_heart_rate. Look for analysis_type to find the type of analysis. "
            "minimum_heart_rate is the lowest heart rate (5 minutes sliding average) registered during the sleep period."
        ),
        (
            "The analysis result for rmssd. Look for analysis_type to find the type of analysis. "
            "rmssd is the average Root Mean Square of Successive Differences (RMSSD) registered during the sleep period."
        ),
        (
            "The analysis result for average_breathing_rate. Look for analysis_type to find the type of analysis. "
            "average_breathing_rate is the average breathing rate registered during the sleep period."
        ),
        (
            "The analysis result for temperature_variation. Look for analysis_type to find the type of analysis. "
            "temperature_variation is the skin temperature deviation from the long-term temperature average."
        ),
    ]
    # False if the output should directly passed back to the planner.
    # True if it should be stored in datapipe
    output_type: bool = True

    def _execute(
        self,
        inputs: List[Any],
    ) -> str:
        df = pd.read_json(
            StringIO(inputs[0]["data"].strip()), orient="records"
        )
        analysis_type = inputs[1].strip()
        if analysis_type == "average":
            df = df.drop("date", axis=1)  # No average for date!
            df = df.mean().to_frame().T
        elif analysis_type == "trend":
            df = self._calculate_slope(df)
        else:
            raise ValueError(
                "The input analysis type has not been defined!"
            )
        df = df.round(2)
        json_out = df.to_json(orient="records")
        return json_out

tasks/google_translator.py
class GoogleTranslate(BaseTask):
    """
    **Description:**

        This task uses google translate to autmatically convert from the user language to english or vise versa.

    """

    name: str = "google_translator"
    chat_name: str = "GoogleTranslator"
    description: str = (
        "Translates queries between different languages."
    )
    dependencies: List[str] = []
    inputs: List[str] = [
        "text to be translated",
        "destination language",
    ]
    outputs: List[str] = []
    output_type: bool = False

    translator: Any = None  #: :meta private:

    @model_validator(mode="before")
    def validate_environment(cls, values: Dict) -> Dict:
        """
            Validate that api key and python package exists in environment.

        Args:
            cls (object): The class itself.
            values (Dict): The dictionary containing the values for validation.
        Return:
            Dict:The original values.
        Raise:
            ImportError: If the 'playwright' package is not installed.

        """

        try:
            from googletrans import Translator

            values["translator"] = Translator()
        except ImportError:
            raise ValueError(
                "Could not import googletrans python package. "
                "Please install it with `pip install googletrans-py`."
            )
        return values

    def _parse_input(
        self,
        input_args: str,
    ) -> List[str]:
        """
            Parse the input string into a list of strings.

        Args:
            input (str): Input string to be parsed.
        Return:
            List[str]: List of parsed strings.

        """
        return input_args.split("$#")

    def _execute(
        self,
        inputs: List[Any] = None,
    ) -> str:
        """
            Abstract method representing the execution of the task.

        Args:
            input (str): Input data for the task.
        Return:
            str: Result of the task execution.
        Raise:
            NotImplementedError: Subclasses must implement the execute method.

        """
        if len(inputs) < 2:
            return "", ""
        dest = inputs[1] if inputs[1] is not None else "en"
        result = self.translator.translate(inputs[0], dest=dest)
        return result.text, result.src

    def explain(
        self,
    ) -> str:
        """
            Provide a sample explanation for the task.

        Return:
            str: Sample explanation for the task.

        """

        return "This task uses google translate to translate between languages"

tasks/serpapi.py
class SerpAPI(BaseTask):
    """
    **Description:**

        This code defines a class named SerpAPI, which is a specific implementation of the abstract BaseTask class.
        The SerpAPI class represents a task that utilizes the SerpAPI (Google Search API) to perform internet searches
        and retrieve relevant information.

    """

    name: str = "serpapi"
    chat_name: str = "InternetSearchSerp"
    description: str = (
        "A low-cost Google Search API."
        "Useful for when you need to answer questions about current events."
    )
    dependencies: List[str] = []
    inputs: List[str] = ["It should be a search query."]
    outputs: List[str] = []
    output_type: bool = False

    search_engine: Any = None  #: :meta private:
    params: Dict = Field(
        default={
            "engine": "google",
            "google_domain": "google.com",
            "gl": "us",
            "hl": "en",
        }
    )
    serpapi_api_key: Optional[str] = None
    aiosession: Optional[aiohttp.ClientSession] = None

    @model_validator(mode="before")
    def validate_environment(cls, values: Dict) -> Dict:
        """
            Validate that api key and python package exists in environment.

        Args:
            values (Dict): The dictionary of attribute values.
        Return:
            Dict: The updated dictionary of attribute values.
        Raise:
            ValueError: If the SerpAPI python package is not installed.

        """

        serpapi_api_key = get_from_dict_or_env(
            values, "serpapi_api_key", "SERPAPI_API_KEY"
        )
        values["serpapi_api_key"] = serpapi_api_key
        try:
            from serpapi import GoogleSearch

            values["search_engine"] = GoogleSearch
        except ImportError:
            raise ValueError(
                "Could not import serpapi python package. "
                "Please install it with `pip install google-search-results`."
            )
        return values

    def get_params(self, query: str) -> Dict[str, str]:
        """
            Get parameters for SerpAPI.

        Args:
            query (str): The search query.
        Return:
            Dict[str, str]: The parameters for the SerpAPI.

        """

        _params = {
            "api_key": self.serpapi_api_key,
            "q": query,
        }
        params = {**self.params, **_params}
        return params

    def results(self, query: str) -> Dict:
        """
            Run query through SerpAPI and return the raw result.

        Args:
            query (str): The search query.
        Return:
            Dict: The raw result from the SerpAPI.

        """

        params = self.get_params(query)
        search = self.search_engine(params)
        res = search.get_dict()
        return res

    @staticmethod
    def _process_response(res: Dict) -> str:
        """
            Process response from SerpAPI.

        Args:
            res (Dict): The raw response from the SerpAPI.
        Return:
            str: Processed information from the SerpAPI response.

        """

        try:
            if "answer_box" in res:
                toret = (
                    "url: "
                    + res["answer_box"]["link"]
                    + "\nmetadata: "
                    + res["answer_box"]["snippet"]
                )
            else:
                toret = (
                    "url: "
                    + res["organic_results"][0]["link"]
                    + "\nmetadata: "
                    + res["organic_results"][0]["snippet"]
                )
        except KeyError:
            return "Could not get the proper response from the search. Try another search query."
        return toret

    def _execute(
        self,
        inputs: List[Any] = None,
    ) -> str:
        """
            Run query through SerpAPI and parse result.

        Args:
            input (str): The input, which should be a search query.
        Return:
            str: The parsed result from the SerpAPI.

        """
        if len(inputs) == 0:
            return ""
        return self._process_response(self.results(inputs[0]))

    def explain(
        self,
    ) -> str:
        """
            Provide an explanation of the task.

        Return:
            str: Explanation of the SerpAPI task.

        """

        return (
            "This task searched in the internet using google search engine, returns the url"
            "and the first top result of the google search."
        )

tasks/affect/sleep_get.py
class SleepGet(Affect):
    """
    **Description:**

        This tasks gets sleep affect data for specific patient.
    """

    name: str = "affect_sleep_get"
    chat_name: str = "AffectSleepGet"
    description: str = (
        "Get the sleep parameters for a specific date or "
        "a period (if two dates are provided). "
        "You must Call $affect_sleep_analysis$ whenever sleep "
        "analysis (e.g., 'average' or 'trend') is needed. DON'T rely on your analysis"
    )
    dependencies: List[str] = []
    inputs: List[str] = [
        "user ID in string. It can be refered as user, patient, individual, etc. Start with 'par_' following with a number (e.g., 'par_1').",
        "start date of the sleep data in string with the following format: '%Y-%m-%d'",
        (
            "end date of the sleep data in string with the following format: '%Y-%m-%d'. "
            "If there is no end date, the value should be an empty string (i.e., '')"
        ),
    ]
    outputs: List[str] = [
        "total_sleep_time (in minutes) is Total amount of sleep (a.k.a. sleep duration) registered during the sleep period.",
        "awake_duration (in minutes) is the total amount of awake time registered during the sleep period.",
        "light_sleep_duration (in minutes) is the total amount of light (N1 or N2) sleep registered during the sleep period.",
        "rem_sleep_duration (in minutes) is the total amount of REM sleep registered during the sleep period.",
        "deep_sleep_duration (in minutes) is the total amount of deep (N3) sleep registered during the sleep period.",
        "sleep_onset_latency (in minutes) is detected latency from bedtime_start to the beginning of the first five minutes of persistent sleep.",
        "midpoint_time_of_sleep (in minutes) is the time from the start of sleep to the midpoint of sleep. The midpoint ignores awake periods.",
        "sleep_efficiency is the percentage of the sleep period spent asleep (100% * sleep duration / time in bed).",
        "average_heart_rate is the average heart rate registered during the sleep period.",
        "minimum_heart_rate is the lowest heart rate (5 minutes sliding average) registered during the sleep period.",
        "rmssd is the average Root Mean Square of Successive Differences (RMSSD) registered during the sleep period.",
        "average_breathing_rate is the average breathing rate registered during the sleep period.",
        "temperature_variation is the skin temperature deviation from the long-term temperature average.",
    ]
    # False if the output should directly passed back to the planner.
    # True if it should be stored in datapipe
    output_type: bool = True
    #
    file_name: str = "sleep.csv"
    device_name: str = "oura"
    local_dir: str = "data/affect"
    columns_to_keep: List[str] = [
        "date",
        "total",
        "awake",
        "light",
        "rem",
        "deep",
        "onset_latency",
        "midpoint_time",
        "efficiency",
        "hr_average",
        "hr_lowest",
        "rmssd",
        "breath_average",
        "temperature_delta",
    ]
    columns_revised: List[str] = [
        "date",
        "total_sleep_time",
        "awake_duration",
        "light_sleep_duration",
        "rem_sleep_duration",
        "deep_sleep_duration",
        "sleep_onset_latency",
        "midpoint_time_of_sleep",
        "sleep_efficiency",
        "average_heart_rate",
        "minimum_heart_rate",
        "rmssd",
        "average_breathing_rate",
        "temperature_variation",
    ]
    variables_in_seconds: List[str] = [
        "total_sleep_time",
        "awake_duration",
        "light_sleep_duration",
        "rem_sleep_duration",
        "deep_sleep_duration",
        "sleep_onset_latency",
        "midpoint_time_of_sleep",
    ]

    def _execute(
        self,
        inputs: List[Any],
    ) -> str:
        user_id = inputs[0].strip()
        full_dir = os.path.join(
            self.local_dir, user_id, self.device_name
        )
        df = self._get_data(
            local_dir=full_dir,
            file_name=self.file_name,
            start_date=inputs[1].strip(),
            end_date=inputs[2].strip(),
            usecols=self.columns_to_keep,
        )
        df.columns = self.columns_revised
        df = self._convert_seconds_to_minutes(
            df, self.variables_in_seconds
        )
        df = df.round(2)
        json_out = df.to_json(orient="records")
        return json_out

tasks/affect/activity_get.py
class ActivityGet(Affect):
    """
    **Description:**

        This tasks gets activity affect data for specific patient.
    """

    name: str = "affect_activity_get"
    chat_name: str = "AffectActivityGet"
    description: str = (
        "Get the physical activity parameters for a specific date or "
        "a period (if two dates are provided). "
        "You must Call $affect_analysis$ whenever physical activity "
        "analysis (e.g., 'average', 'sum', or 'trend') is needed. DON'T rely on your analysis"
    )
    dependencies: List[str] = []
    inputs: List[str] = [
        "user ID in string. It can be refered as user, patient, individual, etc. Start with 'par_' following with a number (e.g., 'par_1').",
        "start date of the physical activity data in string with the following format: '%Y-%m-%d'",
        (
            "end date of the physical activity data in string with the following format: '%Y-%m-%d'."
            "If there is no end date, the value should be an empty string (i.e., '')"
        ),
    ]
    outputs: List[str] = [
        "steps_count is the total number of steps registered during the day.",
        "rest_time is the time (in minutes) during the day spent resting, i.e. sleeping or lying down.",
        "inactive_time is the time (in minutes) during the day spent resting, i.e. sitting or standing still.",
        "low_acitivity_time is the (in minutes) during the day with low intensity activity (e.g. household work).",
        "medimum_acitivity_time is the (in minutes) during the day with medium intensity activity (e.g. walking).",
        "high_acitivity_time is the (in minutes) during the day with high intensity activity (e.g. running).",
    ]

    # False if the output should directly passed back to the planner.
    # True if it should be stored in datapipe
    output_type: bool = True
    #
    file_name: str = "activity.csv"
    device_name: str = "oura"
    local_dir: str = "data/affect"

    columns_to_keep: List[str] = [
        "date",
        "steps",
        "rest",
        "inactive",
        "low",
        "medium",
        "high",
    ]
    columns_revised: List[str] = [
        "date",
        "steps_count",
        "rest_time",
        "inactive_time",
        "low_acitivity_time",
        "medimum_acitivity_time",
        "high_acitivity_time",
    ]

    def _execute(
        self,
        inputs: List[Any] = None,
    ) -> str:
        user_id = inputs[0].strip()
        full_dir = os.path.join(
            self.local_dir, user_id, self.device_name
        )
        df = self._get_data(
            local_dir=full_dir,
            file_name=self.file_name,
            start_date=inputs[1].strip(),
            end_date=inputs[2].strip(),
            usecols=self.columns_to_keep,
        )
        df.columns = self.columns_revised
        df = df.round(2)
        json_out = df.to_json(orient="records")
        return json_out

tasks/test_file.py
class TestFile(BaseTask):
    name: str = "test_file"
    chat_name: str = "TestFile"
    description: str = "analyzes the image and returns description."
    dependencies: List[str] = []
    inputs: List[str] = ["the image file name"]
    outputs: List[str] = []
    output_type: bool = False
    return_direct: bool = True

    translator: Any = None  #: :meta private:

    def parse_input(
        self,
        input: str,
    ) -> List[str]:
        """
        Parse the input string into a list of strings.

        Args:
            input (str): Input string to be parsed.
        Return:
            List[str]: List of parsed strings.

        Example:
            .. code-block:: python

                from langchain import ReActChain, OpenAI
                react = ReAct(llm=OpenAI())

        """

        return input.split("$#")

    def execute(
        self,
        input: str,
    ) -> str:
        """
        Abstract method representing the execution of the task.

        Args:
            input (str): Input data for the task.
        Return:
            str: Result of the task execution.
        Raise:
            NotImplementedError: Subclasses must implement the execute method.

        Example:
            .. code-block:: python

                from langchain import ReActChain, OpenAI
                react = ReAct(llm=OpenAI())

        """

        self.parse_input(input)
        return "this image is a classification results of a data"

    def explain(
        self,
    ) -> str:
        """
        Provide a sample explanation for the task.

        Return:
            str: Sample explanation for the task.

        Example:
            .. code-block:: python

                from langchain import ReActChain, OpenAI
                react = ReAct(llm=OpenAI())

        """

        return "This task simply asks user to provide more information or continue interaction."

tasks/read_from_datapipe.py
class ReadDataPipe(BaseTask):
    """
    **Description:**

      This code reads raw data stored in datapipe. When different tasks are executed, there are situations that the final data is stored
      in the datapipe when the final called task's output_type=True. In these situations, this task is called to retireve the latest stored data
      to be used for final inference.
    """

    name: str = "read_from_datapipe"
    chat_name: str = "DataPipeReader"
    description: str = (
        "Get the stored information from datapipe to be used to answer user query accurately. "
        "This should be called when the final answer is in datapipe."
    )
    dependencies: List[str] = []
    inputs: List[str] = [
        "the datapipe key in the format $datapipe:key$"
    ]
    outputs: List[str] = []
    output_type: bool = False

    def _execute(
        self,
        inputs: List[Any] = None,
    ) -> str:
        """
            This simply retrieves data from datapipe.

        Args:
            inputs (List[Any]): The datapipe key
        Return:
            str: The raw data along with the instructions.

        """
        if len(inputs) == 0:
            return ""
        return (
            "The data along with the description for each data is provided. "
            "Use the data and description to provide a detailed answer regarding the user query.\n\n"
            + json.dumps(inputs[0])
        )

    def explain(
        self,
    ) -> str:
        """
            Provide an explanation of the task.

        Return:
            str: Explanation of the SerpAPI task.

        """
        return "This task is to read data from datapipe."

tasks/affect/activity_analysis.py
class ActivityAnalysis(Affect):
    """
    **Description:**

        This tasks performs average, sum, or trend analysis on the provided raw activity affect data for specific patient.
    """

    name: str = "affect_activity_analysis"
    chat_name: str = "AffectActivityAnalysis"
    description: str = (
        "Analyze the physical activity data. You must Call this whenever physical activity analysis"
        "(e.g., 'average', 'sum', or 'trend') is needed. DON'T rely on your analysis."
        "For example, if the user asks for trends (or variations) in data, you must call this task"
    )
    dependencies: List[str] = ["affect_activity_get"]
    inputs: List[str] = [
        "It is an string but in json format. It is the output of the $affect_activity_get$",
        "analysis_type. It can be one of [$average$, $sum$, $trend$].",
    ]
    outputs: List[str] = [
        (
            "The analysis result for steps_count. Look for analysis_type to find the type of analysis. "
            "steps_count is the total number of steps registered during the day."
        ),
        (
            "The analysis result for rest_time. Look for analysis_type to find the type of analysis. "
            "rest_time is the time (in minutes) during the day spent resting, i.e. sleeping or lying down."
        ),
        (
            "The analysis result for inactive_time. Look for analysis_type to find the type of analysis. "
            "inactive_time is the time (in minutes) during the day spent resting, i.e. sitting or standing still."
        ),
        (
            "The analysis result for low_acitivity_time. Look for analysis_type to find the type of analysis. "
            "low_acitivity_time is the (in minutes) during the day with low intensity activity (e.g. household work)."
        ),
        (
            "The analysis result for medimum_acitivity_time. Look for analysis_type to find the type of analysis. "
            "medimum_acitivity_time is the (in minutes) during the day with medium intensity activity (e.g. walking)."
        ),
        (
            "The analysis result for high_acitivity_time. Look for analysis_type to find the type of analysis. "
            "high_acitivity_time is the (in minutes) during the day with high intensity activity (e.g. running)."
        ),
    ]
    # False if the output should directly passed back to the planner.
    # True if it should be stored in datapipe
    output_type: bool = False

    def _execute(
        self,
        inputs: List[Any] = None,
    ) -> str:
        if len(inputs) == 0:
            return ""

        df = pd.read_json(
            StringIO(inputs[0]["data"].strip()), orient="records"
        )
        analysis_type = inputs[1].strip()
        if analysis_type == "average":
            df = df.drop("date", axis=1)  # No average for date!
            df = df.mean().to_frame().T
        elif analysis_type == "sum":
            df = df.drop("date", axis=1)  # No sum for date!
            df = df.sum().to_frame().T
        elif analysis_type == "trend":
            df = self._calculate_slope(df)
        else:
            raise ValueError(
                "The input analysis type has not been defined!"
            )
        df = df.round(2)
        json_out = df.to_json(orient="records")
        return json_out

tasks/task_types.py
class TaskType(str, Enum):
    SERPAPI = "serpapi"
    CLICK = "click"
    GET_CURRENT_PAGE = "current_page"
    EXTRACT_HYPERLINKS = "extract_hyperlinks"
    EXTRACT_TEXT = "extract_text"
    GET_ELEMENTS = "get_elements"
    NAVIGATE_BACK = "navigate_back"
    NAVIGATE = "navigate"
    AFFECT_SLEEP_GET = "affect_sleep_get"
    AFFECT_ACTIVITY_GET = "affect_activity_get"
    AFFECT_SLEEP_ANALYSIS = "affect_sleep_analysis"
    AFFECT_ACTIVITY_ANALYSIS = "affect_activity_analysis"
    GOOGLE_TRANSLATE = "google_translate"
    ASK_USER = "ask_user"
    READ_FROM_DATAPIPE = "read_from_datapipe"
    TEST_FILE = "test_file"

Answer: TaskType.GOOGLE_TRANSLATE: GoogleTranslate,

<think>
The dictionary TASK_TO_CLASS is being built to map each TaskType enum to its corresponding class. The snippet ends after the ActivityAnalysis entry with a trailing comma, indicating more entries follow. The next imported class after the imports is GoogleTranslate, so the next mapping should associate TaskType.GOOGLE_TRANSLATE with GoogleTranslate.
</think>
TaskType.GOOGLE_TRANSLATE: GoogleTranslate,