<file_sep>import torch
from collections import defaultdict
from itertools import chain
import os


class TransferTransfoConstants:
    SPECIAL_TOKENS = ["<bos>", "<eos>", "<speaker1>", "<speaker2>", "<pad>"]
    MODEL_INPUTS = ["input_ids", "mc_token_ids", "lm_labels", "mc_labels", "token_type_ids"]
    PADDED_INPUTS = ["input_ids", "lm_labels", "token_type_ids"]

    ATTR_TO_SPECIAL_TOKEN = {
        'bos_token': '<bos>',
        'eos_token': '<eos>',
        'pad_token': '<pad>',
        'additional_special_tokens': ["<speaker1>", "<speaker2>", "<end>", "<eot>"]
    }

    ADDITIONAL_TOKENS = ["_nofact"]

class TransferTransfoWithKnowledgeConstants(object):
    SPECIAL_TOKENS = TransferTransfoConstants.SPECIAL_TOKENS
    MODEL_INPUTS = TransferTransfoConstants.MODEL_INPUTS
    PADDED_INPUTS = TransferTransfoConstants.PADDED_INPUTS

    ADDITIONAL_TOKENS = TransferTransfoConstants.ADDITIONAL_TOKENS + [
        "<person>", "</person>",
        "<genre>", "</genre>",
        "<movie_title>", "</movie_title>",
    ]

    ATTR_TO_SPECIAL_TOKEN = TransferTransfoConstants.ATTR_TO_SPECIAL_TOKEN


def collate_batch_elements(batch, tokenizer, device, pad_left=False):

    batch_inputs = defaultdict(list)

    for instance in batch:
        for field, data in instance.items():
            batch_inputs[field].append(data)
    
    pad_token_map = {
        "labels": -100,
        "input_ids": tokenizer.eos_token_id,
        "default": 0
    }
    padded_inputs = ["input_ids", "labels", "token_type_ids"]
    
    model_inputs = ["input_ids", "labels", "token_type_ids"]

    if pad_left:
        model_inputs.append("attention_mask")
    padded_batch = pad_batch_items(batch_inputs, pad_token_map, padded_inputs, pad_left)
    tensorized_input = []

    for input_name in model_inputs:
        tensor = torch.tensor(padded_batch[input_name], device=device)
        tensorized_input.append(tensor)
    
    return tensorized_input


def collate_transfertransfo_batch_elements(batch, tokenizer, args):
    batch_inputs = defaultdict(list)
    chained_batch = chain(*batch)

    for instance in chained_batch:
        for field, data in instance.items():
            batch_inputs[field].append(data)


    pad_token = tokenizer.convert_tokens_to_ids(TransferTransfoConstants.SPECIAL_TOKENS[-2])

    padded_dataset = pad_dataset(batch_inputs, padding=pad_token)

    tensorized_input = []

    batch_size = tuple([len(batch_inputs[TransferTransfoConstants.MODEL_INPUTS[0]]) // args.num_candidates])
    for input_name in TransferTransfoConstants.MODEL_INPUTS:
        tensor = torch.tensor(padded_dataset[input_name])

        if input_name != "mc_labels":
            tensor = tensor.view((-1, args.num_candidates) + tensor.shape[1:])
        else:
            tensor = torch.ones(size=batch_size, dtype=torch.long) * (args.num_candidates - 1)
        tensorized_input.append(tensor)
    return tensorized_input


def pad_dataset(dataset, padding=0):
    """ Pad the dataset. This could be optimized by defining a Dataset class and padding at the batch level, but this is simpler. """
    max_l = max(len(x) for x in dataset["input_ids"])
    for name in TransferTransfoConstants.PADDED_INPUTS:
        dataset[name] = [x + [padding if name != "lm_labels" else -100] * (max_l - len(x)) for x in dataset[name]]
    return dataset


def pad_batch_items(batch_items, pad_token_map, padded_inputs, pad_left):
    max_seq_len = max(len(x) for x in batch_items["input_ids"])

    default_pad_token = pad_token_map["default"]

    if pad_left:
        # Attention mask is necessary to avoid attending on left padding tokens
        # this isn't a problem for the right-padding case since
        # the logits from the right padding tokens can be ignored.
        # See: https://github.com/huggingface/transformers/issues/808
        batch_items["attention_mask"] = [[0 if i < max_seq_len - len(x) else 1 for i in range(max_seq_len)] for x in batch_items["input_ids"]]

    for name in padded_inputs:
        pad_token = pad_token_map.get(name, default_pad_token)

        if pad_left:
            # Experimenting with left padding for batch inference
            batch_items[name] = [ ([pad_token] * (max_seq_len - len(x)) + x) for x in batch_items[name]]
        else:
            batch_items[name] = [ (x + [pad_token] * (max_seq_len - len(x))) for x in batch_items[name]]
    return batch_items        


def save_model_config_and_tokenizer(config, tokenizer, args):
    config_path = os.path.join(args.experiment_path, args.experiment_name)
    os.makedirs(config_path, exist_ok=True)
    config.save_pretrained(config_path)
    tokenizer.save_pretrained(config_path)


def save_model_checkpoint(model, args, checkpoint_name="checkpoint.pt"):
    checkpoint_path = os.path.join(args.experiment_path, args.experiment_name, "checkpoints")
    os.makedirs(checkpoint_path, exist_ok=True)
    checkpoint_file_path = os.path.join(checkpoint_path, checkpoint_name)
    torch.save(model.state_dict(), checkpoint_file_path)


def save_full_model(model, tokenizer, args, model_name):
    checkpoint_path = os.path.join(args.experiment_path, args.experiment_name, model_name)
    
    model.save_pretrained(checkpoint_path)
    tokenizer.save_pretrained(checkpoint_path)

<file_sep>import json
import csv

def load_conversations(conversations_file):
    with open(conversations_file, 'r') as conversations_file:
        conversations = []

        for line in conversations_file:
            conversations.append(json.loads(line.strip()))
    
        return conversations

def dump_conversations_to_file(conversations, output_file):
    with open(output_file, 'w') as conversations_file:
        for conversation in conversations:
            conversations_file.write(json.dumps(conversation))
            conversations_file.write('\n')

def load_imdb_list(file_path):
    with open(file_path, 'r') as imdb_file:
        entities = [row["name"] for row in csv.DictReader(imdb_file)]

        return entities

def popular_actors_list():
    actors_list = load_imdb_list('top_1000_actors.csv')

    actors_list.extend([
        "<NAME>",
        "<NAME>", # Tesla in The Prestige
        "<NAME>", # 
        "<NAME>", # Joan from Mad Men
        "<NAME>", # Dr. Who
        "<NAME>", # Finn from Force Awakens
        "<NAME>", # Rey from Force Awakens
        "<NAME>", # SNL
        "<NAME>", # Dude from Scrubs
        "<NAME>", # Surprised this bad mofo isn't on the list
        "<NAME>", # Masters of Sex
        "<NAME>", # The iconic dracula
        "<NAME>", # Mr. Bean, Johnny English
        "<NAME>", # Baby Driver
        "<NAME>", # High School Musical
        "<NAME>", # Dinesh from Silicon Valley
        "<NAME>", # 8 Mile
        "<NAME>", # The Simpsons
        "<NAME>", #The Best Little Whorehouse in Texas
        "<NAME>", # Troy from Community
        "<NAME>", 
        "<NAME>", # Whole bunch of martial arts movies
        "<NAME>", # Grey's Anatomy
        "<NAME>", # Eraserhead and Twin Peaks
        "<NAME>", # Watson from Sherlock ; Shaun of the Dead
        "<NAME>", # Honey I shrunk .. ; Ghostbusters
        "<NAME>", # Girls Trip
        "<NAME>", # The Breakfast Club 
        "<NAME>", # Hamilton
        "<NAME>", # <NAME>, Deadpool
        "<NAME>", # Maleficent
        "<NAME>", # <NAME> from a nightmare on Elm Street
        "<NAME>", # Ong Bak
        "<NAME>", # Criminal Minds
    ])
    return actors_list

def popular_directors_list():
    directors_list =  load_imdb_list('top_250_directors.csv')

    directors_list.extend([
        "<NAME>", # Grand Budapest Hotel, Isle of Dogs etc.
        "<NAME>", # A nightmare on Elm Street
        "<NAME>", # Star wars
        "<NAME>", # Baby Driver
        "<NAME>", # Any movie with Johnny Depp / <NAME>
        "<NAME>", # Paprika
        "<NAME>", "<NAME>", # South Park, Orgazmo etc.
        "<NAME>", # Iron Man, Avengers
        "<NAME>", # Pacific Rim, The Shape of Water
        "<NAME>", # BlacKKKlansman, Da Five Bloods
        "<NAME>", # Escape from New York/LA
        "<NAME>", # Nosferatu
        "<NAME>", # Star Trek/Star Wars the Force Awakens
        "<NAME>", # <NAME>
        "<NAME>", # Harry Potter
        "<NAME>", # Crouching Tiger, Hidden Dragon; Brokeback Mountain
        "<NAME>", # Get Out
        "<NAME>", # Doctor Sleep, The Haunting of Hill House, Hush
        "<NAME>", # Star Wars the Last Jedi
        "<NAME>", # Black Swan, Requiem for a Dream
        "<NAME>", # Man of Steel, Watchmen, Justice League
        "<NAME>", # Platoon, Snowden
    ])

    return directors_list<file_sep>import random

from torch.utils.data import Dataset

from itertools import chain

class RedialDialoGPTDataset(Dataset):
    """
    This dataset class prepares input in the vanilla DialoGPT format:
    
    <|endoftext|> TURN_1 <|endoftext|> TURN_2 <|endoftext|> ... <|endoftext|> TURN_N <|endoftext|>

    This baseline is expected to perform worse since there's 
    nothing for the model to condition on
    """

    def __init__(self, dataset, tokenizer, args):
        self.dataset = dataset
        self.tokenizer = tokenizer
        self.max_history_turns = args.max_history_turns
        self.inference = args.inference

    def __getitem__(self, index):

        example = self.dataset[index]

        example = self._truncate_example(example)
        instance = self.prepare_input_for_dialogpt(example, self.tokenizer, self.inference)

        return instance

    def _truncate_example(self, example):

        if len(example.context) > self.max_history_turns:
            truncated_context = example.context[-self.max_history_turns:]
            example = example._replace(context=truncated_context)

        return example

    @staticmethod
    def prepare_input_for_dialogpt(example, tokenizer, inference=False):
        
        bos, eos = tokenizer.convert_tokens_to_ids([tokenizer.bos_token, tokenizer.eos_token])
        context_turns = example.context
        system_turn = example.response
        
        input_ids = [bos]
        labels = [-100]
        
        for turn in context_turns:
            input_ids += turn + [eos]
            labels += len(turn) * [eos]
        
        if not inference: # Don't add system response for inference
            input_ids += system_turn + [eos]
            labels += system_turn + [eos]
        
        token_type_ids = [0 for _ in labels]
        instance = {
            "input_ids": input_ids,
            "labels": labels,
            "token_type_ids": token_type_ids
        }

        return instance


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



class RedialTransferTransfoDataset(Dataset):
    """
    Mimics the setup from the TransferTransfo paper. This setup has been
    fairly reliable for knowledge grounded models such as Topical Chats.
    """

    def __init__(self, dataset, tokenizer, special_tokens, args):
        self.dataset = dataset
        self.tokenizer = tokenizer
        self.special_tokens = special_tokens

        self.dataset = dataset
        self.tokenizer = tokenizer
        self.special_tokens = special_tokens

        self.max_history = args.max_history_turns
        self.num_candidates = args.num_candidates

    @staticmethod
    def sample_candidates(dataset, index, num_samples):
        # Sample candidates IID
        candidates = [response for (_, response, _) in random.sample(dataset, num_samples)]

        return candidates

    @staticmethod
    def build_input_from_segments(context, response, knowledge, tokenizer, special_tokens, lm_labels):

        #print("context:", context)
        #print("response:", response)
        #print("knowledge:", knowledge)

        bos, eos, speaker1, speaker2 = tokenizer.convert_tokens_to_ids(special_tokens[:4])

        sequence = [[bos] + knowledge] + context + [response + [eos]]
        sequence = [sequence[0]] + [
            [(speaker2 if (len(sequence) - i) % 2 else speaker1)] + s for i, s in enumerate(sequence[1:])]

        instance = {}
        instance["input_ids"] = list(chain.from_iterable(sequence))
        instance["token_type_ids"] = [speaker2 if (len(sequence) - i - 1) % 2 else speaker1 for i, s in enumerate(sequence) for _ in s]
        instance["mc_token_ids"] = len(instance["input_ids"]) - 1

        if lm_labels:
            instance["lm_labels"] = ([-100] * sum(len(s) for s in sequence[:-1])) + [-100] + sequence[-1][1:]
        else:
            instance["lm_labels"] = [-100] * len(instance["input_ids"])

        #print(tokenizer.decode(instance["input_ids"]))
        #print(instance["mc_token_ids"])
        #input(">>>")

        return instance

    @staticmethod
    def truncate_inputs(context, max_context_exchanges):
        # Limit the number of context exchanges to retain
        if len(context) > (2 * max_context_exchanges + 1):
            context = context[-(2 * max_context_exchanges + 1):]

        return context

    def __getitem__(self, index):

        context, response, knowledge = self.dataset[index]

        context = self.truncate_inputs(context, self.max_history)
        tokenized_context = [self.tokenizer.encode(turn) for turn in context]
        tokenized_response = self.tokenizer.encode(response)

        candidates = self.sample_candidates(self.dataset, index, self.num_candidates - 1)

        tokenized_candidates = [self.tokenizer.encode(candidate) for candidate in candidates]
        tokenized_candidates.append(tokenized_response)

        tokenized_knowledge = []

        for (type, item) in knowledge:
            #print(f" * type: {type}\titem: {item}")
            if type == "dact":
                tokenized_knowledge.extend(self.tokenizer.encode(item))
            else:
                start_tag = f"<{type}>"
                end_tag = f"</{type}>"
                tokenized_knowledge.append(self.tokenizer.convert_tokens_to_ids(start_tag))
                tokenized_knowledge.extend(self.tokenizer.encode(item))
                tokenized_knowledge.append(self.tokenizer.convert_tokens_to_ids(end_tag))
        instances = []

        for j, candidate in enumerate(tokenized_candidates):
            lm_labels = bool(j == self.num_candidates - 1)
            instance = self.build_input_from_segments(tokenized_context, candidate, tokenized_knowledge, self.tokenizer, self.special_tokens, lm_labels)
            instances.append(instance)

        return instances

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

    def get_num_batches(self, batch_size):
        return len(self) // batch_size # has the effect of performing "drop last"<file_sep># ReDial Data

More info at https://github.com/RaymondLi0/conversational-recommendations .
The repository contains the code for NeurIPS 2018 paper "Towards Deep Conversational Recommendations" https://arxiv.org/abs/1812.07617

# imdb info
More files can be found in the shared drive at CSE243 > DATA . 
- `retrieved-imdb-info.zip`: imdb info on movies talked about
- `traintest-with-movie-info.zip`: imdb movie info aligned with redial dialogue turns.
<file_sep>
class Metric(object):
    def get(self):
        raise NotImplementedError("Subclasses must implement this")

class RunningMetric(Metric):
    def __init__(self):
        self.current_value = 0.0
        self.num_steps = 0

    def _add(self, value):
        self.current_value += (value - self.current_value) / (self.num_steps + 1)
        self.num_steps += 1

    def add(self, value):
        self._add(value)

    def get(self):
        return self.current_value


class RunningLambdaMetric(RunningMetric):
    def __init__(self, fn):
        self.func = fn
        super().__init__()

    def add(self, *args, **kwargs):
        self._add(float(self.func(*args, **kwargs)))

class MetricLambda(Metric):
    def __init__(self, fn, metric):
        self.func = fn
        self.metric = metric

    def get(self):
        return self.func(self.metric.get())

class Accuracy(RunningLambdaMetric):
    def __init__(self):
        super().__init__(self._compute_accuracy)

    def _compute_accuracy(self, logits, labels):
        predictions = logits.argmax(dim=-1)

        return (predictions == labels).mean()
<file_sep># conversational-movie-recommender<file_sep>import re

class DBPedia(object):
    # These correspond to labels that have been misrecognized
    BLACKLIST_URIS = {"http://dbpedia.org/resource/Glossary_of_tennis_terms", 
    "http://dbpedia.org/resource/Good_Movie",
    "http://dbpedia.org/resource/Sierra_Entertainment",
    "http://dbpedia.org/resource/Nice",
    "http://dbpedia.org/resource/Take_Care_(album)",
    "http://dbpedia.org/resource/Cloning",
    "http://dbpedia.org/resource/Blood",
    "http://dbpedia.org/resource/Downhill_creep",
    "http://dbpedia.org/resource/Movies",
    "http://dbpedia.org/resource/Hey_There",
    "http://dbpedia.org/resource/Swimming_(sport)",
    "http://dbpedia.org/resource/Princess_Falls",
    "http://dbpedia.org/resource/Haha_(entertainer)",
    "http://dbpedia.org/resource/LOL",
    "http://dbpedia.org/resource/Drag_queen",
    "http://dbpedia.org/resource/Yea_Football_Club",
    "http://dbpedia.org/resource/Oh_Yeah_(Yello_song)",
    "http://dbpedia.org/resource/Scalable_Coherent_Interface",
    "http://dbpedia.org/resource/CAN_bus",
    "http://dbpedia.org/resource/The_New_One_(horse)",
    "http://dbpedia.org/resource/Information_technology",
    "http://dbpedia.org/resource/The_Glad_Products_Company",
    "http://dbpedia.org/resource/AM_broadcasting",
    "http://dbpedia.org/resource/To_Heart",
    "http://dbpedia.org/resource/National_Organization_for_Women",
    "http://dbpedia.org/resource/Hit_or_Miss_(New_Found_Glory_song)",
    "http://dbpedia.org/resource/Canada",
    "http://dbpedia.org/resource/Different_Things",
    "http://dbpedia.org/resource/Norwegian_Trekking_Association",
    "http://dbpedia.org/resource/Take_One_(Canadian_magazine)",
    "http://dbpedia.org/resource/For_Inspiration_and_Recognition_of_Science_and_Technology",
    "http://dbpedia.org/resource/Two_Guys",
    "http://dbpedia.org/resource/The_Sydney_Morning_Herald",
    "http://dbpedia.org/resource/Booting",
    "http://dbpedia.org/resource/Precious_Time_(album)",
    "http://dbpedia.org/resource/I\\u0027m_Glad",
    "http://dbpedia.org/resource/Social_Democratic_Party_of_Switzerland",
    "http://dbpedia.org/resource/International_Maritime_Organization",
    "http://dbpedia.org/resource/LOL",
    "http://dbpedia.org/resource/Names_of_God_in_Judaism",
    "http://dbpedia.org/resource/Ike_Turner",
    "http://dbpedia.org/resource/Tricky_Stewart",
    "http://dbpedia.org/resource/Movies!",
    }

multi_spaces_pattern = re.compile(r"\s+")

def process_text(text):
    return multi_spaces_pattern.sub(" ", text.capitalize())<file_sep>import argparse
import logging
import math
from pprint import pformat

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import GPT2Tokenizer, AdamW

from dataset_utils import (
    prepare_redial_baseline_dataset,
    prepare_redial_knowledge_grounded_dataset,
    get_movie_db_map
)
from datasets import (RedialDialoGPTDataset, RedialTransferTransfoDataset)
from train_utils import (
    TransferTransfoConstants,
    TransferTransfoWithKnowledgeConstants,
    collate_batch_elements,
    collate_transfertransfo_batch_elements,
    save_model_checkpoint,
    save_full_model,
    save_model_config_and_tokenizer)
from trainer.counter import GlobalStepCounter
from trainer.metrics import RunningMetric, MetricLambda, RunningLambdaMetric
from trainer.scheduler import PiecewiseLinearLR

logger = logging.getLogger(__file__)


def prepare_dataloaders(args, tokenizer):

    movie_db_map = get_movie_db_map(args.movies_data_path)
    special_terms = None

    if args.configuration == "baseline":
        dataset = prepare_redial_baseline_dataset(
            args.data_path,
            tokenizer,
            movie_db_map,
            args.data_cache_path
            )

        train_dataset, test_dataset = RedialDialoGPTDataset(dataset["train"], tokenizer, args), \
            RedialDialoGPTDataset(dataset["test"], tokenizer, args)

        def collate_fn(batch):
            return collate_batch_elements(batch, tokenizer, args.device)

        train_loader, test_loader = \
            DataLoader(
                train_dataset,
                batch_size=args.train_batch_size,
                collate_fn=collate_fn,
                shuffle=True), \
            DataLoader(
                test_dataset,
                batch_size=args.test_batch_size,
                collate_fn=collate_fn,
                shuffle=False)
    else:
        dataset, special_terms = prepare_redial_knowledge_grounded_dataset(
            args.data_path,
            tokenizer,
            movie_db_map,
            args.data_cache_path,
            split_files={"train": args.train_file, "test": args.eval_file},
            recommender_only=args.recommender_only,
            include_dacts=args.include_dialog_acts
        )
        special_terms.extend([
        "<cast>", "</cast>",
        "<movie_genre>", "</movie_genre>",
        "<director>", "</director>",
        ])
        train_dataset = RedialTransferTransfoDataset(dataset["train"], tokenizer, TransferTransfoConstants.SPECIAL_TOKENS, args)
        test_dataset = RedialTransferTransfoDataset(dataset["test"], tokenizer, TransferTransfoConstants.SPECIAL_TOKENS, args)

        def collate_fn(batch):
            return collate_transfertransfo_batch_elements(batch, tokenizer, args)

        train_loader = DataLoader(train_dataset, batch_size=args.train_batch_size,
                                  collate_fn=collate_fn,
                                  shuffle=True)
        test_loader = DataLoader(test_dataset, batch_size=args.test_batch_size,
                                  collate_fn=collate_fn,
                                  shuffle=False)

    return train_loader, test_loader, special_terms


def train_lm(model, loader, optimizer, scheduler, step_counter, args):
    logger.info("Running training")
    model.train()

    running_loss = RunningMetric()

    for i, batch in enumerate(tqdm(loader)):
        input_ids, labels, token_type_ids = batch

        # Forward pass and loss computation
        lm_loss, *_ = model(input_ids=input_ids, labels=labels, token_type_ids=token_type_ids)
        # lm_loss is given by sum(batch_loss) / batch_size

        # We need to average loss over all gradient accumulation steps
        loss = lm_loss / args.gradient_accumulation_steps

        running_loss.add(loss.item())

        # Backprop step. Pytorch automatically accumulates
        # gradients each time backward is called
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)

        # Parameter update step
        if (i + 1) % args.gradient_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()

        if (i + 1) % args.log_every_n == 0:
            logger.info(f"Iteration {i + 1}: [Running Loss: {running_loss.get()};Running PPL: {math.exp(running_loss.get())}]")


        step_counter.step()

        if step_counter.get() % args.checkpoint_every_n == 0:
            checkpoint_name = f"checkpoint-{step_counter.get()}.pt"
            save_model_checkpoint(model, args, checkpoint_name)
            logger.info(f"Model checkpoint {checkpoint_name} saved!")

    logger.info(f"Training loss: {running_loss.get()}")
    logger.info(f"Training PPL: {math.exp(running_loss.get())}")


def train_double_heads_lm(model, loader, optimizer, scheduler, step_counter, args):
    running_loss = RunningMetric()
    ppl = MetricLambda(math.exp, running_loss)

    for i, batch in enumerate(tqdm(train_loader)):
        model.train()
        batch = tuple(input_tensor.to(args.device) for input_tensor in batch)

        input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch

        (lm_loss), (mc_loss), *_ = model(
            input_ids, token_type_ids=token_type_ids, mc_token_ids=mc_token_ids,
            mc_labels=mc_labels, lm_labels=lm_labels
        )

        loss = (lm_loss * args.lm_coef + mc_loss * args.mc_coef) / args.gradient_accumulation_steps
        running_loss.add(float(loss))
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)

        if (i + 1) % args.log_every_n == 0:
            logger.info(
                f"Iteration {i}: [Running Loss: {running_loss.get()};Running PPL: {ppl.get()}]")

        if i % args.gradient_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()



        scheduler.step()
        step_counter.step()

        if step_counter.get() % args.checkpoint_every_n == 0:
            checkpoint_name = f"checkpoint-{step_counter.get()}.pt"
            save_model_checkpoint(model, args, checkpoint_name)
            logger.info(f"Model checkpoint {checkpoint_name} saved!")

    logger.info(f"Epoch loss: {running_loss.get()}")
    logger.info(f"Epoch PPL: {ppl.get()}")


def evaluate_lm(model, loader, loss_fn, args):
    logger.info("Running evaluation")
    model.eval()
    running_loss = RunningMetric()
    with torch.no_grad():
        for i, batch in enumerate(tqdm(loader)):
            input_ids, labels, token_type_ids = batch

            # No loss is returned if lm_labels is not passed
            lm_logits, *_ = model(
                input_ids=input_ids, token_type_ids=token_type_ids
            )

            lm_logits_flattened = lm_logits.view(-1, lm_logits.size(-1))
            labels_flattened = labels.view(-1)
            loss = loss_fn(lm_logits_flattened, labels_flattened)
            running_loss.add(loss.item())

            if (i + 1) % args.log_every_n == 0:
                logger.info(f"Iteration {i}: [Running Loss: {running_loss.get()};Running PPL: {math.exp(running_loss.get())}]")

    logger.info(f"Validation NLL: {running_loss.get()}")
    logger.info(f"Validation PPL: {math.exp(running_loss.get())}")


def evaluate_double_heads_lm(model, val_loader, loss_fn, args):
    model.eval()

    running_nll = RunningLambdaMetric(loss_fn)
    ppl = MetricLambda(math.exp, running_nll)

    with torch.no_grad():
        for i, batch in enumerate(tqdm(val_loader)):
            batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
            input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch

            lm_logits, mc_logits, *_ = model(
                input_ids, token_type_ids=token_type_ids, mc_token_ids=mc_token_ids,
            )

            lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(-1, lm_logits.size(-1))
            lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)

            running_nll.add(lm_logits_flat_shifted, lm_labels_flat_shifted)

    logger.info(f"NLL Loss: {running_nll.get()}")
    logger.info(f"Perlexity: {ppl.get()}")


def train_baseline_lm(model, loaders, optimizer, loss_fn, scheduler, args):

    train_loader, test_loader = loaders

    step_counter = GlobalStepCounter()
    for i in range(args.n_epochs):
        logger.info(f"Epoch {i + 1}:")
        train_lm(model, train_loader, optimizer, scheduler, step_counter, args)
        evaluate_lm(model, test_loader, loss_fn, args)
        epoch_model = f"{args.experiment_name}_epoch_{i + 1}"
        save_full_model(model, tokenizer, args, epoch_model)
        logger.info(f"Model {epoch_model} saved!")
        logger.info(f"Epoch {i + 1} completed!\n")

    if args.n_epochs < 1:
        evaluate_lm(model, test_loader, optimizer, args)


def train_knowledge_grounded_lm(model, loaders, optimizer, loss_fn, scheduler, args):
    train_loader, val_loader = loaders

    step_counter = GlobalStepCounter()
    for epoch in range(args.n_epochs):

        train_double_heads_lm(model, train_loader, optimizer, scheduler, step_counter, args)
        epoch_model = f"{args.experiment_name}_epoch_{epoch + 1}"
        save_full_model(model, tokenizer, args, epoch_model)
        logger.info(f"Model {epoch_model} saved!")
        logger.info(f"Epoch {epoch + 1} completed!\n")

        evaluate_double_heads_lm(model, val_loader, loss_fn, args)


    if args.n_epochs < 1:
        evaluate_double_heads_lm(model, val_loader, loss_fn, args)


def setup_baseline_lm(args):
    if args.gradient_checkpoint:
        from models.gpt2 import GPT2LMHeadModel
        model = GPT2LMHeadModel.from_pretrained(args.model_checkpoint)
    else:
        from transformers import GPT2LMHeadModel
        model = GPT2LMHeadModel.from_pretrained(args.model_checkpoint)

    return model


def setup_double_heads_model(new_vocabulary_size, args):
    if args.gradient_checkpoint:
        from models.gpt2 import GPT2DoubleHeadsModel
        model = GPT2DoubleHeadsModel.from_pretrained(args.model_checkpoint)
    else:
        from transformers import GPT2DoubleHeadsModel
        model = GPT2DoubleHeadsModel.from_pretrained(args.model_checkpoint)

    model.resize_token_embeddings(new_num_tokens=new_vocabulary_size)
    return model


def setup_knowledge_grounded_tokenizer(args):
    tokenizer = GPT2Tokenizer.from_pretrained(args.model_checkpoint)

    logger.info("Adding special tokens to tokenizer:")
    logger.info(TransferTransfoWithKnowledgeConstants.ADDITIONAL_TOKENS + list(TransferTransfoWithKnowledgeConstants.ATTR_TO_SPECIAL_TOKEN.items()))
    # Add special tokens
    num_added_norm_tokens = tokenizer.add_tokens(TransferTransfoWithKnowledgeConstants.ADDITIONAL_TOKENS)
    num_added_special_tokens = tokenizer.add_special_tokens(TransferTransfoWithKnowledgeConstants.ATTR_TO_SPECIAL_TOKEN)

    new_vocab_size = len(tokenizer.encoder) + num_added_norm_tokens + num_added_special_tokens
    # Return tokenizer and new vocabulary size

    logger.info("Successfully added: %s tokens", num_added_norm_tokens + num_added_special_tokens)
    return tokenizer, new_vocab_size

def get_argparser():
    parser = argparse.ArgumentParser()

    parser.add_argument('--configuration',
        default='baseline',
        choices=['baseline', 'knowledge_grounded']
    )

    parser.add_argument('--model_checkpoint',
        default="microsoft/DialoGPT-medium",
        help="The model checkpoint to use"
    )

    parser.add_argument('--gradient_checkpoint',
        action="store_true",
        help="Use gradient checkpoint variant GPT2"
    )

    parser.add_argument('--data_path',
        default="redial/",
        help="Path to dataset"
    )
    parser.add_argument('--movies_data_path',
        default="redial/movies_with_mentions.csv",
        help="Path to movie mentions file"
    )
    parser.add_argument('--train-file',
        default="train_data_swda_tagged.jsonl",
        help="Name of train jsonl file."
    )
    parser.add_argument('--eval-file',
        default="test_data_swda_tagged.jsonl",
        help="Name of test jsonl file."
    )
    parser.add_argument('--data_cache_path',
        default="redial_dataset_cache.pkl",
        help="Path to cached data"
    )
    parser.add_argument('--experiment_path',
        default="runs",
        help="Parent directory for experiments"
    )
    parser.add_argument('--experiment_name',
        required=True,
        help="Name of the experiment"
    )
    parser.add_argument('--seed',
        default=42,
        type=int,
        help="Random seed for experiments"
    )

    parser.add_argument('--max_history_turns',
        default=2,
        type=int,
        help="How many turns from context to retain"
    )

    parser.add_argument('--n_epochs',
        default=3,
        type=int,
        help="Number of epochs for training"
    )

    parser.add_argument('--train_batch_size',
        default=4,
        type=int,
        help="Train batch size"
    )
    parser.add_argument('--test_batch_size',
        default=4,
        type=int,
        help="Test batch size"
    )
    parser.add_argument('--lr',
        default=6.25e-5,
        type=float,
        help="Base/initial learning rate"
    )
    parser.add_argument('--device',
        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
        help="Device to store the model"
    )
    parser.add_argument('--gradient_accumulation_steps',
        default=16,
        type=int,
        help="Number of steps to accumulate gradient for"
    )
    parser.add_argument('--max_norm',
        default=1.0,
        type=float,
        help="Maximum norm for the gradient"
    )
    parser.add_argument('--log_every_n',
        default=500,
        type=int,
        help="The frequency (in number of steps) with which information is shown"
    )
    parser.add_argument('--checkpoint_every_n',
        default=1800,
        type=int,
        help="The frequency (in number of steps) with which the model checkpoints are saved"
    )
    parser.add_argument('--recommender_only',
        dest='recommender_only',
        action='store_true',
        help="Train only on recommender side utterances"
    )

    parser.set_defaults(include_dialog_acts=True)
    parser.add_argument('--exclude_dialog_acts',
                        dest='include_dialog_acts',
                        action='store_false',
                        help="Whether to exclude dialog act in the knowledge")

    # Double heads model specific args
    double_heads_parser = parser.add_argument_group('Double Heads Model Arguments:')
    double_heads_parser.add_argument('--num_candidates',
                                     type=int, default=2,
                                     help="Number of candidates to select from")
    double_heads_parser.add_argument('--lm_coef',
                                     type=float, default=1.0,
                                     help="Weight for Language Model loss coefficient")
    double_heads_parser.add_argument('--mc_coef',
                                     type=float, default=1.0,
                                     help="Weight for Multiple-Choice loss coefficient")


    parser.set_defaults(recommender_only=False)
    return parser

if __name__ == "__main__":
    parser = get_argparser()

    args = parser.parse_args()
    args.inference = False
    logger.info(f"Arguments : {pformat(args)}")

    logging.basicConfig(level=logging.INFO)
    logger.info("Arguments: %s", pformat(args))

    logger.info("Prepare tokenizer, pretrained model and optimizer.")

    logger.info("Load datasets")

    if args.configuration == "baseline":
        tokenizer = GPT2Tokenizer.from_pretrained(args.model_checkpoint)
        model = setup_baseline_lm(args)
    else:
        tokenizer, new_vocab_size = setup_knowledge_grounded_tokenizer(args)
        model = setup_double_heads_model(new_vocab_size, args)

    train_loader, test_loader, special_terms = prepare_dataloaders(args, tokenizer)
    if isinstance(special_terms, list):
        print("adding more tokens:")
        print(" *", special_terms)
        num_added = tokenizer.add_tokens(special_terms)
        print("num_added:", num_added)
        new_vocab_size += num_added

    if args.configuration == "baseline":
        model = setup_baseline_lm(args)
    else:
        model = setup_double_heads_model(new_vocab_size, args)


    optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
    scheduler = PiecewiseLinearLR(optimizer, [(0, args.lr), (args.n_epochs * len(train_loader), 0.0)])
    loss_fn = nn.CrossEntropyLoss(ignore_index=-100)  # used for computing validation loss
    model.to(args.device)


    save_model_config_and_tokenizer(model.config, tokenizer, args)

    if args.configuration == "baseline":
        train_baseline_lm(model, (train_loader, test_loader), optimizer, loss_fn, scheduler, args)
    else:
        train_knowledge_grounded_lm(model, (train_loader, test_loader), optimizer, loss_fn, scheduler, args)



<file_sep>"""
 Created by diesel
 12/23/19
"""

from collections import defaultdict

import nltk
import json

import spacy
nlp = spacy.load('en_core_web_sm')

class Lexicon(object):
    _save_attrs = ["name", "terms", "_id2term", "info", "_id_counter",
                   "_standard_2_id", "_all_refs"]
    def __init__(self, name):
        self.name = name
        self.terms = []
        self._id2term = {}
        self.info = {}
        self._id_counter = 0
        self._standard_2_id = {}
        self._all_refs = defaultdict(list)
        self._phrases = None


    def _clean_up(self):

        for record in self._id2term.values():
            record["referential_forms"] = list(set(record["referential_forms"]))


    def to_dict(self):
        self._clean_up()

        d = {}
        for k in self._save_attrs:
            d[k] = getattr(self, k)
        return d


    def to_lines(self , delim="; ", num_spaces=2):
        self._clean_up()

        all_lines = []
        for entry in self.terms:

            info = self.info[entry["id"]]

            category = self.name
            if category == "characters":
                category = "character"

            category = category.upper()

            lines = [category, "standard_form" + ": "+ entry["standard_form"]]
            forms = entry.get("referential_forms", [])
            if forms:
                forms = delim.join(forms)
                lines.append("forms: " + forms)

            for attr, val in info.items():
                if attr in {"forms", "full_name", "standard_form"}:
                    continue

                if isinstance(val, list):
                    val = delim.join(val)
                lines.append(attr + ": " + str(val))

            lines.extend([""] * num_spaces)
            all_lines.extend(lines)

        return all_lines


    def to_df(self , delim="; ", num_spaces=2):
        self._clean_up()

        all_lines = []
        for entry in self.terms:
            info = self.info[entry["id"]]
            category = self.name
            if category == "characters":
                category = "character"
            category = category.upper()
            e = {
                "category": category,
                "standard_form": entry["standard_form"],
                "forms": json.dumps(entry.get("referential_forms", []))
            }


            for attr, val in info.items():
                if attr in {"forms", "full_name", "standard_form"}:
                    continue
                if isinstance(val, list):
                    val = json.dumps(val)
                e[attr] = val



            all_lines.append(e)

        return all_lines




    @classmethod
    def from_dict(cls, d):
        assert "name" in d
        lex = cls(d["name"])
        for k in cls._save_attrs:
            if d.get(k):
                setattr(lex, k, d[k])

        lex._id2term = {}
        _terms = []
        if len(lex.terms) > 0:
            for t in lex.terms:
                if t["id"] not in lex._id2term:
                    lex._id2term[t["id"]] = t
                    _terms.append(t)

        lex.terms = _terms
        assert len(lex.terms) == len(lex._id2term), "{} {}".format(len(lex.terms), len(lex._id2term))

        return lex



    def get_info(self, tid_or_standard_form, default_val=None):
        info = self.info.get(tid_or_standard_form)
        if info is None:
            key = self._standard_2_id.get(tid_or_standard_form)
            info = self.info.get(key)

        if info is None:
            info = default_val
        return info


    def new_id(self):
        self._id_counter += 1
        return self._id_counter


    def new_term(self, standard_form, info=None, referential_forms=None, tid=None, standard_form_is_id=True):


        if standard_form in self._standard_2_id:
            the_id = self._standard_2_id[standard_form]
            for ref in referential_forms:
                self._add_new_ref(ref, the_id)

        else:
            new_term = {
                "standard_form": standard_form,
                "referential_forms": referential_forms if referential_forms else [],
                "id": (standard_form if standard_form_is_id else self.new_id()) if tid is None else tid
            }
            self.terms.append(new_term)
            self._id2term[new_term["id"]] = new_term
            self.info[new_term["id"]] = info

            for ref in referential_forms:
                self._add_new_ref(ref, new_term["id"])


            self._standard_2_id[standard_form] = new_term["id"]


    def update_info(self, tid, new_info):
        tinfo = self.info.get(tid, {})
        tinfo.update(new_info)


    def _add_new_ref(self, ref, tid):
        if ref not in self._all_refs or tid not in self._all_refs[ref]:
            if ref not in self._all_refs:
                self._all_refs[ref] = []
            self._all_refs[ref].append(tid)



    def new_referential_phrase(self, refs, standard_form=None, tid=None):
        assert standard_form or tid

        if tid is None:
            tid = standard_form

        if not isinstance(refs, list):
            refs = [refs]

        #print("tid:", [tid])
        record = self._id2term.get(tid)
        #print("record:", record)
        record["referential_forms"].extend([r for r in refs if r not in record["referential_forms"]])
        for ref in refs:
            self._add_new_ref(ref, tid)
            print(" * add {} => {}".format(ref, tid))






    @property
    def phrases(self):
        if self._phrases is None:
            self._phrases = set(self._all_refs.keys())

        return self._phrases


    def ref2standard(self, ref, get_tid=False):

        ids = self._all_refs.get(ref)
        if ids:
            if not isinstance(ids, list):
                ids = [ids]
            standard_forms = [self._id2term[_i].get("standard_form") for _i in ids]

        else:
            standard_forms = None

        retval = standard_forms, ids if get_tid else standard_forms
        return retval




class LexBuilder(object):
    def __init__(self):
        self._entries = defaultdict(list)
        self._lexicons = None

    @property
    def lexicons(self):
        return self._lexicons

    def add_entry(self, e, etype=None):
        if etype is None:
            etype = e["category"]
        self._entries[etype].append(e)

    def get_entry_key_set(self):
        return set([k for etype, ents in self._entries.items() for e in ents for k in e.keys()])

    def lex_as_dict(self):
        return {lex.name: lex.to_dict() for lex in self.lexicons.values()}


    def build_lexes(self):
        lexicons = {}
        for ent_type, entries in self._entries.items():
            lex = Lexicon(ent_type)
            for entry in entries:
                if isinstance(entry, str):
                    entry = entries[entry]
                if not entry.get("full_name") or entry["full_name"] == "none none":
                    continue
                refs = self.process_ref(entry["forms"])
                lex.new_term(entry["full_name"], info=entry, referential_forms=refs)
                #refs = [process_ref(player["full_name"])]
                #player_lex.new_term(player["full_name"], info=player, referential_forms=refs)
            lexicons[ent_type] = lex
        self._lexicons = lexicons


    def save_lexes(self, outfile):
        with open(outfile, "w") as fout:
            json.dump(self.lex_as_dict(), fout, indent=2)

    @staticmethod
    def process_ref(text):
        def _process(ref):
            doc = nlp(ref)
            toks = " ".join(getattr(tok, "text", None) for tok in doc).lower()
            #toks = " ".join(nltk.word_tokenize(ref)).lower()
            return toks
        retval = [_process(ref) for ref in text] if isinstance(text, list) else _process(text)
        return retval




def main():
    pass


if __name__ == "__main__":
    main()
<file_sep>import argparse
import torch

from tqdm.auto import tqdm
from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPT2Config

import torch.nn.functional as F
from torch.utils.data import DataLoader
from datasets import RedialDialoGPTDataset, RedialTransferTransfoDataset
from dataset_utils import get_movie_db_map, prepare_redial_baseline_dataset, prepare_redial_knowledge_grounded_dataset
from train_utils import collate_batch_elements, collate_transfertransfo_batch_elements, TransferTransfoConstants

def prepare_dataloader(args, tokenizer):

    movie_db_map = get_movie_db_map(args.movies_data_path)
    dataset = prepare_redial_baseline_dataset(
        args.data_path,
        tokenizer,
        movie_db_map,
        args.data_cache_path
        )

    test_dataset = RedialDialoGPTDataset(dataset["test"], tokenizer, args)


    collate_fn = lambda batch: collate_batch_elements(batch, tokenizer, args.device, pad_left=True)

    test_loader = DataLoader(test_dataset, batch_size=args.test_batch_size, collate_fn=collate_fn, shuffle=False)

    return test_loader

def prepare_knowledge_grounded_dataloader(args, tokenizer):
    movie_db_map = get_movie_db_map(args.movies_data_path)
    dataset, special_terms = prepare_redial_knowledge_grounded_dataset(
        args.data_path,
        tokenizer,
        movie_db_map,
        args.data_cache_path,
        split_files={"train": args.train_file, "test": args.eval_file},
        recommender_only=args.recommender_only,
        include_dacts=args.include_dialog_acts
    )
    special_terms.extend([
        "<cast>", "</cast>",
        "<movie_genre>", "</movie_genre>",
        "<director>", "</director>",
    ])

    test_dataset = RedialTransferTransfoDataset(
        dataset["test"], tokenizer, TransferTransfoConstants.SPECIAL_TOKENS, args)

    def collate_fn(batch):
        return collate_transfertransfo_batch_elements(batch, tokenizer, args)

    test_loader = DataLoader(test_dataset, batch_size=args.test_batch_size, collate_fn=collate_fn, shuffle=False)

    return test_loader


def save_outputs(outputs, args):
    with open(args.output_file_path, 'w') as output_file:
        output_file.write("\n".join(outputs))


def top_filtering(logits, top_k=0., top_p=0.9, threshold=-float('Inf'), filter_value=-float('Inf')):
    """ Filter a distribution of logits using top-k, top-p (nucleus) and/or threshold filtering
        Args:
            logits: logits distribution shape (vocabulary size)
            top_k: <=0: no filtering, >0: keep only top k tokens with highest probability.
            top_p: <=0.0: no filtering, >0.0: keep only a subset S of candidates, where S is the smallest subset
                whose total probability mass is greater than or equal to the threshold top_p.
                In practice, we select the highest probability tokens whose cumulative probability mass exceeds
                the threshold top_p.
            threshold: a minimal threshold to keep logits
    """
    assert logits.dim() == 1  # Only work for batch size 1 for now - could update but it would obfuscate a bit the code

    top_k = min(top_k, logits.size(-1))
    if top_k > 0:
        # Remove all tokens with a probability less than the last token in the top-k tokens
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value

    if top_p > 0.0:
        # Compute cumulative probabilities of sorted tokens
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probabilities = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # Remove tokens with cumulative probability above the threshold
        sorted_indices_to_remove = cumulative_probabilities > top_p
        # Shift the indices to the right to keep also the first token above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # Back to unsorted indices and set them to -infinity
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[indices_to_remove] = filter_value

    indices_to_remove = logits < threshold
    logits[indices_to_remove] = filter_value

    return logits


def decode_sequences(input_ids, token_type_ids, model, tokenizer, args):

    special_tokens_ids = tokenizer.convert_tokens_to_ids(TransferTransfoConstants.SPECIAL_TOKENS)

    outputs = []
    for i in range(len(input_ids)):
        input_seq = tokenizer.decode(input_ids[i][0])
        prefix, suffix = input_seq.rsplit("<speaker", maxsplit=1)
        context = prefix + "<speaker" + suffix[:2]  # Hacky way to append the speaker tag
        current_output = []

        attempts = 0
        # Keep trying to generate output until a limited number of times
        expanded_tok_type_ids = token_type_ids[i][0].tolist()
        for j in range(args.max_length):  # Add trailing tokens
            expanded_tok_type_ids.append(expanded_tok_type_ids[-1])
        expanded_tok_type_ids = torch.tensor(expanded_tok_type_ids).to(args.device)
        for j in range(args.max_length):
            prefix_input_seq = torch.tensor(tokenizer.encode(context) + current_output).unsqueeze(0)
            truncated_tok_type_ids = expanded_tok_type_ids[:prefix_input_seq.shape[-1]].unsqueeze(0)
            logits = model(prefix_input_seq.to(args.device), token_type_ids=truncated_tok_type_ids.to(args.device))

            if isinstance(logits, tuple) or len(logits.shape) == 4:  # for gpt2 and maybe others
                logits = logits[0]
            logits = logits[0, -1, :] / args.temperature
            logits = top_filtering(logits, top_k=args.top_k, top_p=args.top_p)
            probs = F.softmax(logits, dim=-1)

            prev = torch.topk(probs, 1)[1] if args.no_sample else torch.multinomial(probs, 1)
            if prev.item() in special_tokens_ids:
                patience = 3
                while prev.item() in special_tokens_ids:
                    if probs.max().item() == 1 or patience == 0:
                        # Disabled this rather noisy warning
                        # logger.warn("Warning: model generating special token with probability 1.")
                        break  # avoid infinitely looping over special token
                    prev = torch.multinomial(probs, num_samples=1)
                    patience -= 1
            if prev.item() in special_tokens_ids:
                break
            current_output.append(prev.item())

        output = tokenizer.decode(current_output)
        outputs.append(output.replace('\n', ''))

    return outputs


def generate_outputs(model, loader, tokenizer, args):

    
    all_outputs = []
    with torch.no_grad():
        for i, batch in enumerate(tqdm(loader)):

            input_ids, _, _, attention_mask = batch
            output_sequences = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                do_sample=not args.no_sample,
                pad_token_id=tokenizer.eos_token_id,
                max_length=input_ids.size(1) + args.max_length,
                temperature=args.temperature,
                top_k=args.top_k,
                top_p=args.top_p
                )

            for x in output_sequences:
                all_outputs.append(tokenizer.decode(x[input_ids.size(1):], skip_special_tokens=True))
    
    save_outputs(all_outputs, args)

def generate_outputs_single_example(model, loader, tokenizer, args):
    outputs = []
    with torch.no_grad():
        for i, batch in enumerate(tqdm(loader)):
            input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
            outputs += decode_sequences(input_ids, token_type_ids, model, tokenizer, args)

    save_outputs(outputs, args)

def main(args):
    tokenizer = GPT2Tokenizer.from_pretrained(args.model_configuration)

    if args.configuration == "baseline":
        test_loader = prepare_dataloader(args, tokenizer)
    else:
        test_loader = prepare_knowledge_grounded_dataloader(args, tokenizer)
    
    if args.model_configuration != args.model_checkpoint:
        config = GPT2Config.from_pretrained(args.model_configuration)
        state_dict = torch.load(args.model_checkpoint)
        # I didnt' read the documentation carefully enough
        model = GPT2LMHeadModel.from_pretrained(None, config=config,state_dict=state_dict)
    else:
        model = GPT2LMHeadModel.from_pretrained(args.model_checkpoint)
    model.to(args.device)

    if args.configuration == "baseline":
        generate_outputs(model, test_loader, tokenizer, args)
    else:
        generate_outputs_single_example(model, test_loader, tokenizer, args)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument('--configuration',
                        type=str,
                        default='baseline',
                        choices=['baseline', 'knowledge_grounded'])

    parser.add_argument('--model_configuration',
        default="microsoft/DialoGPT-medium",
        help="The model configuration to use"
    )

    parser.add_argument('--model_checkpoint',
        default="microsoft/DialoGPT-medium",
        help="The model checkpoint to use"
    )

    parser.add_argument('--data_path',
        default="redial/",
        help="Path to dataset"
    )

    parser.add_argument('--movies_data_path',
        default="redial/movies_with_mentions.csv",
        help="Path to movie mentions file"
    )

    parser.add_argument('--train-file',
                        default="train_data_swda_tagged.jsonl",
                        help="Name of train jsonl file."
                        )
    parser.add_argument('--eval-file',
                        default="test_data_swda_tagged.jsonl",
                        help="Name of test jsonl file."
                        )

    parser.add_argument('--data_cache_path',
        default="redial_dataset_cache.pkl",
        help="Path to cached data"
    )

    parser.add_argument('--test_batch_size',
        default=1,
        type=int,
        help="Test batch size"
    )

    parser.add_argument('--model_metadata_path', type=str, default='./runs/bert_swbd_pd_nrg',
                        help='Path to the tokenizer and model configuration')
    parser.add_argument('--output_file_path', type=str, default='submissions/submissions.txt')

    parser.add_argument('--max_history_turns',
        default=2,
        type=int,
        help="How many turns from context to retain"
    )

    parser.add_argument('--device',
        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
        help="Device to store the model"
    )

    parser.add_argument("--temperature", type=int, default=0.7, help="Sampling softmax temperature")
    parser.add_argument("--top_k", type=int,
                        default=0.,
                        help="Filter top-k tokens before sampling (<=0: no filtering)")
    parser.add_argument("--top_p", type=float, default=0.9,
                        help="Nucleus filtering (top-p) before sampling (<=0.0: no filtering)")
    parser.add_argument("--no_sample", action='store_true', help="Set to use greedy decoding instead of sampling")
    parser.add_argument("--max_length", type=int, default=50, help="Maximum length of the output utterances")

    double_heads_parser = parser.add_argument_group('Double Heads Model Arguments:')
    double_heads_parser.add_argument('--num_candidates',
                                     type=int, default=1,
                                     help="Number of candidates to select from")
    parser.add_argument('--recommender_only',
                        dest='recommender_only',
                        action='store_true',
                        help="Train only on recommender side utterances"
                        )
    parser.set_defaults(include_dialog_acts=True)
    parser.add_argument('--exclude_dialog_acts',
                        dest='include_dialog_acts',
                        action='store_false',
                        help="Whether to exclude dialog act in the knowledge")


    args = parser.parse_args()

    args.inference = True
    print(args)
    main(args)<file_sep>"""
 Created by diesel
 12/23/19
"""

#import scaffold
#print(vars(scaffold))



from lexicon import Lexicon
import json
import nltk
from collections import Counter


"""

        # Penn Tagset
        # Optional DT: (\S+/DT\s*)?
        # Adjetive can be JJ,JJR,JJS: (\S+/JJ\w?\s*)*
        # Noun can be NN,NNS,NNP,NNPS: (\S+/NN\w*\s*)+
        # ?P<> named group
        # ?: non-capturing group
        # See https://docs.python.org/3/howto/regex.html#non-capturing-and-named-groups

        # index/word/pos/ner

        # capital determiner 'A' or 'The' followed by nouns
        # The man
        # (The Markell family) always
        generic = "((?:(?:[0-9]+)/(?:The|A)/DT/\S+\s)(?:\S+/NN\S*/\S+\s*)+)"

        # possesive relations people
        #  our son; his girlfriend
        generic2 = "((?:(?:[0-9]+)/\S+/PRP\$/\S+\s)(?:\S+/NN\S*/\S+\s*)+)"

        # Mister Deals saw
        proper = "((?:^(?:[0-9]+)/\S+/NN\S*/\S+\s)(?:\S+/NN\S*/\S+)*)"

        person = "((?:(?:[0-9]+)/\S+/NN\S*/\S+\s+)*(?:\S+/NN\S*/PERSON\s*))"

        speaker = "((?:[0-9]+/SPEAKER/\S+/\S+\s+))"

        #char_patterns = "|".join([generic, proper, person, speaker])

        re_patterns = {
            "animate_NN_wn": generic,
            "proper_NN": proper,
            "person_NER": person,
            "speaker": speaker,
            "poss_rel_wn": generic2
        }


        #pattern = re.compile(r'(?P<CHAR>(?:\S+/DT\s*)?(?:\S+/JJ\w?\s*)*(?:\S+/NN\w*\s*)+)')
        #pattern = re.compile("(?P<CHAR>{})".format(char_patterns))

        regex = {k: re.compile("(?P<CHAR>{})".format(char_patterns)) for k, char_patterns in re_patterns.items()}

        #found_entities = defaultdict(list)
        for pat_name, pattern in regex.items():
            #print("pattern_name:", pat_name)
            for m in pattern.finditer(tokens):
                mention = [feats.split("/") for feats in m.group("CHAR").rstrip().split()]
"""



class RuleBasedNER():
    def __init__(self, lexicons):
        self._lexicons = lexicons

    def tokenize_text(self, text):
        return [t.lower() for t in nltk.word_tokenize(text)]

    def tag_tokens(self, tokens):
        #print("tokens:", tokens)

        indexed_grams = {}
        for n in [4, 3, 2, 1]:
            ngrams = nltk.ngrams(tokens, n)

            for j, gram in enumerate(ngrams):
                g = " ".join(gram)
                if g in indexed_grams:
                    if not isinstance(indexed_grams[g], list):
                        indexed_grams[g] = [indexed_grams[g]]
                    indexed_grams[g].append(j)
                else:
                    indexed_grams[g] = j

        #print("all_grams:", list(indexed_grams.keys()))

        mentions = []
        _mention_idxs = set()
        for lex_name, lex in self._lexicons.items():

            #print("phrases:", lex.phrases)

            found = set(indexed_grams.keys()) & lex.phrases
            #print("found:", found)
            found = sorted([(len(f),f) for f in found], reverse=True)

            if len(found) > 0:
                for length, gram in found:
                    #print("gram:", gram)
                    idxs = indexed_grams.get(gram)
                    #print("idxs:", idxs)
                    if not isinstance(idxs, list):
                        idxs = [idxs]

                    for start_idx in idxs:
                        #print("start_idx:", start_idx)
                        #print("len:", len(gram.split()))
                        end_idx = start_idx + len(gram.split())

                        span_idxs = set(range(start_idx, end_idx))
                        #print("span_idxs:", span_idxs)
                        #print("start,end:", start_idx, end_idx)
                        standard_form, ent_id = lex.ref2standard(gram, get_tid=True)
                        #print("_mention_idxs:", _mention_idxs)
                        if len(_mention_idxs & span_idxs) < 1:
                            mention = {
                                "ent_id": ent_id,
                                "standard_form": standard_form,
                                "words": tokens[start_idx:end_idx],
                                "start_idx": start_idx,
                                "end_idx": end_idx,
                                "ent_type": lex.name
                            }
                            mentions.append(mention)
                            for _idx_ in span_idxs:
                                _mention_idxs.add(_idx_)

        #print("mentions:", mentions)
        return mentions



def player_analysis():
    dm = DataManager(
        season_dir_path="../data/season_info",
        players_path="../data/nba-static-data/nba-players.json",
        teams_path="../data/nba-static-data/nba-teams.json",
        fun_facts_path="../data/ready-facts.json",
        questions_path="../data/ready-questions.json",
        standings_path="../data/nba-playoff-picture",
    )

    print("All Players:")
    full_names = [player["full_name"] for player in dm.players]
    print("num_names:", len(full_names))
    print("num unique names:", len(set(full_names)))

    dups = []
    _all = set()
    for name in full_names:
        if name in _all:
            dups.append(name)
        _all.add(name)

    for d in dups:
        print(" * ", d)

    print("Active Players:")
    full_names = [player["full_name"] for player in dm.players if player["is_active"]]
    print("num_names:", len(full_names))
    print("num unique names:", len(set(full_names)))


def process_ref(text):
    return " ".join(nltk.word_tokenize(text)).lower()


def setup_lexicons():
    dm = DataManager(
        season_dir_path="../data/season_info",
        players_path="../data/nba-static-data/nba-players.json",
        teams_path="../data/nba-static-data/nba-teams.json",
        fun_facts_path="../data/ready-facts.json",
        questions_path="../data/ready-questions.json",
        standings_path="../data/nba-playoff-picture",
    )


    # set up lexicons of know entities
    player_lex = Lexicon("player")
    for player in dm.players:
        refs = [process_ref(player["full_name"])]
        player_lex.new_term(player["full_name"], info=player, referential_forms=refs)

    team_lex = Lexicon("team")
    for team in dm.teams:
        refs = [process_ref(team[k]) for k in ["full_name", "city", "nickname", ]] # "abbreviation"
        team_lex.new_term(team["full_name"], info=team, referential_forms=refs)

    lexicons = {
        lex.name: lex.to_dict() for lex in [player_lex, team_lex]
    }

    with open("../data/lexicons/lex.json", "w") as fout:
        json.dump(lexicons, fout, indent=2)

    with open("../data/lexicons/lex.json", "r") as fin:
        lexicons2 = json.load(fin)

    lex2 = {
        name: Lexicon.from_dict(d) for name, d in lexicons2.items()
    }

    prepared_lexicons = "../data/lexicons/lex.json"
    with open(prepared_lexicons, "r") as fin:
        lexicons2 = json.load(fin)

    lexicons = {
        name: Lexicon.from_dict(d) for name, d in lexicons2.items()
    }


def team_ner_demo(fact, ner_tagger, results, lex):
    print(("-" * 20) + "\n\n")
    text = fact["text"]
    team = fact["team"]
    toks = [t.lower() for t in nltk.word_tokenize(text)]
    mentions = ner_tagger.tag_tokens(toks)
    #player_mentions = [men for men in mentions if men["ent_type"] == "player"]
    team_mentions = [men for men in mentions if men["ent_type"] == "team"]

    nicknames = [lex.get_info(men["standard_form"][0])["nickname"] for men in team_mentions]
    nicknames = "; ".join(nicknames)

    print("\ntext:", text)
    print("fact team:", team)
    print("team_mention:", "; ".join([" ".join(men["words"]) for men in team_mentions]))
    print("nickname:", nicknames)

    if len(team_mentions) == 1:
        words = " ".join(team_mentions[0]["words"])
        standard = team_mentions[0]["standard_form"][0]
    else:
        words = "; ".join([" ".join(men["words"]) for men in team_mentions])
        standard = "; ".join([men["standard_form"][0] for men in team_mentions])
    #print("player_mention:   ", words)
    print("{} == {}".format(team, nicknames))
    if team.lower() == nicknames.lower():
        result = "PASS"
    else:
        result = "FAIL"
    print(result)
    results.append(result)

def player_ner_demo(fact, ner_tagger, results):
    print(("-" * 20) + "\n\n")
    text = fact["text"]
    player = fact["player"]
    toks = [t.lower() for t in nltk.word_tokenize(text)]
    mentions = ner_tagger.tag_tokens(toks)
    player_mentions = [men for men in mentions if men["ent_type"] == "player"]
    team_mentions = [men for men in mentions if men["ent_type"] == "team"]
    print("\ntext:", text)
    print("player:", player)
    print("team_mention:", "; ".join([" ".join(men["words"]) for men in team_mentions]))
    if len(player_mentions) == 1:
        words = " ".join(player_mentions[0]["words"])
        standard = player_mentions[0]["standard_form"][0]
    else:
        words = "; ".join([" ".join(men["words"]) for men in player_mentions])
        standard = "; ".join([men["standard_form"][0] for men in player_mentions])
    print("player_mention:   ", words)
    print("{} == {}".format(player, standard))
    if player.lower() == standard.lower():
        result = "PASS"
    else:
        result = "FAIL"
    print(result)
    results.append(result)


def load_lexicons(prepared_lexicons="../data/lexicons/lex.json"):

    with open(prepared_lexicons, "r") as fin:
        lexicons = json.load(fin)

    lexicons = {
        name: Lexicon.from_dict(d) for name, d in lexicons.items()
    }

    return lexicons



def main():
    dm = DataManager(
        season_dir_path="../data/season_info",
        players_path="../data/nba-static-data/nba-players.json",
        teams_path="../data/nba-static-data/nba-teams.json",
        fun_facts_path="../data/ready-facts.json",
        questions_path="../data/ready-questions.json",
        standings_path="../data/nba-playoff-picture",
        templates_path="../data/season-templates/templates.json"
    )


    aka_list = [
        {"full_name": "<NAME>", "aka": ["<NAME>", "<NAME>"]},
        {"full_name": "<NAME>", "aka": ["Shaq"]}
    ]


    player_lex = Lexicon("player")
    for player in dm.players:
        refs = [process_ref(player["full_name"])]
        player_lex.new_term(player["full_name"], info=player, referential_forms=refs)

    for aka in aka_list:
        player_lex.new_referential_phrase(
            [process_ref(ref) for ref in aka["aka"]],
            aka["full_name"])


    team_lex = Lexicon("team")
    for team in dm.teams:
        refs = [process_ref(team[k]) for k in ["full_name", "city", "nickname"]] # "abbreviation"
        team_lex.new_term(team["full_name"], info=team, referential_forms=refs)

    lexicons = {
        lex.name: lex for lex in [player_lex, team_lex]
    }

    ner_tagger = RuleBasedNER(lexicons)

    results = []

    for fact in dm.fun_facts["fun-facts"]["players"]["facts"]:
        player_ner_demo(fact, ner_tagger, results)

    results = Counter(results)
    print("results:")
    for k,v in results.items():
        print("{}: {}".format(k, v))
    print("accuracy:", results["PASS"]/ (results["PASS"] + results["FAIL"]))


    results = []

    for fact in dm.fun_facts["fun-facts"]["teams"]["facts"]:
        team_ner_demo(fact, ner_tagger, results, lexicons["team"])

    results = Counter(results)
    print("results:")
    for k,v in results.items():
        print("{}: {}".format(k, v))
    print("accuracy:", results["PASS"]/ (results["PASS"] + results["FAIL"]))

    prepared_lexicons = "../data/lexicons/lex.json"
    lexicons = {name: lex.to_dict() for name, lex in lexicons.items()}
    with open(prepared_lexicons, "w") as fout:
        json.dump(lexicons, fout, indent=2)


if __name__ == "__main__":
    main()
<file_sep>import csv

import requests
import json
import logging
import argparse
import os
import spacy
import imdb
import pdb

from data_loader import load_ner_tagger
from data_utils import load_conversations, dump_conversations_to_file, popular_actors_list, popular_directors_list

from annotation_utils import DBPedia, process_text

from fuzzywuzzy import fuzz, process

import pdb
from tqdm.auto import tqdm

from collections import defaultdict

nlp = spacy.load("en_core_web_lg")

def spotlight_annotate(text):
    """
    Annotate the spotlight text with the default arguments
    """

    payload = {
        "text": text
    }
    headers = {
        "Accept": "application/json"
    }

    response = requests.get(
        "http://localhost:2222/rest/annotate", 
        headers=headers,
        params=payload
        )

    # print(response.text)
    return response.json()

def spotlight_annotate_conversations(conversations):

    categories_to_not_stop_at = {
        "http://dbpedia.org/resource/Romantic_comedy_film",
        "http://dbpedia.org/resource/Cartoon",
        "http://dbpedia.org/resource/Romance_film",
        "http://dbpedia.org/resource/Thriller_(genre)",
        "Lol! i am not a big tom fan at all."
    }

    # TODO : Use a secondary NER to identify mentions more accurately
    blacklist_types = {"Schema:MusicAlbum"}

    recognized_category_types = defaultdict(int)

    for conversation in tqdm(conversations):

        for message in conversation["messages"]:

            if not message["text"]:
                continue
            spotlight_annotations = spotlight_annotate(process_text(message["text"]))

            cleaned_mentions = []
            
            if spotlight_annotations.get("Resources"):
                for entity in spotlight_annotations["Resources"]:
                    if entity["@URI"] in DBPedia.BLACKLIST_URIS:
                        break
                    if entity["@types"]:
                        for t in entity["@types"].split(","):
                            recognized_category_types[t] += 1
                # else:
                    # if entity["@URI"] not in categories_to_not_stop_at:                   
                    #     pdb.set_trace()
                cleaned_mentions = [entity for entity in spotlight_annotations["Resources"] if entity["@URI"] not in DBPedia.BLACKLIST_URIS]
            
            message["spotlight_mentions"] = cleaned_mentions
    return conversations
    



def spacy_ner_annotate(conversations):
    recognized_category_types = defaultdict(int)

    for conversation in tqdm(conversations):
        for message in conversation["messages"]:
            doc = nlp(process_text(message["text"]))

            message_entities = []
            for ent in doc.ents:
                message_entities.append({
                    "surface": ent.text,
                    "type": ent.label_ 
                })

                recognized_category_types[ent.label_] += 1
            
            message["spacy_mentions"] = message_entities
    
    print(recognized_category_types)

    return conversations


def link_mention_to_imdb(conversations, imdb_sqlite_path=None):
    top_actors = popular_actors_list()
    top_directors = popular_directors_list()

    if imdb_sqlite_path:
        ia = imdb.IMDb('s3', os.path.join('sqlite+pysqlite:///', args.imdb_sqlite_path))
    else:
        ia = imdb.IMDb()

    match_cache = {}

    for conversation in tqdm(conversations):
        for message in conversation["messages"]:
            imdb_records = []

            if "spotlight_mentions" in message:

                for mention in message["spotlight_mentions"]:
                    if "Schema:Person" in mention["@types"]:
                        # Actor disambiguation logic
                        surface_form = mention["@surfaceForm"]

                        if surface_form == "Goodnight":
                            continue
                        
                        if surface_form.lower() == "jim" and "jim carrey" in message["text"].lower():
                            print("Jim converted to jim carrey")
                            surface_form = "Jim Carrey"

                        top_director_match, top_director_score = process.extractOne(surface_form, top_directors)

                        if top_director_score == 100:
                            # Tolerate only exact matches for director
                            print("Matched director ", top_director_match, "for ", surface_form)

                            
                            cached_result = match_cache.get(best_match)

                            if cached_result:
                                imdb_records.append(cached_result)
                            else:
                                imdb_person_results = ia.search_person(best_match)

                                if len(imdb_person_results) > 0:
                                    imdb_records.append(imdb_person_results[0])
                                    match_cache[best_match] = imdb_person_results[0]
                        else:
                            top_actor_matches = process.extract(surface_form, top_actors, limit=5, scorer=fuzz.partial_ratio)

                            num_high_matches = len([match for match, score in top_actor_matches if score > 80])

                            if num_high_matches > 1:
                                print("Warning! Multiple matches for ", surface_form, " ; Top matches: ", top_actor_matches)
                                num_100_matches = len([match for match, score in top_actor_matches if score == 100])

                                if num_100_matches == 1:
                                    
                                    best_match = top_actor_matches[0][0]
                                    print("Choosing closest match ", best_match)

                                    cached_result = match_cache.get(best_match)

                                    if cached_result:
                                        imdb_records.append(cached_result)
                                    else:
                                        imdb_person_results = ia.search_person(best_match)
                                        if len(imdb_person_results) > 0:
                                            imdb_records.append(imdb_person_results[0])
                                            match_cache[best_match] = imdb_person_results[0]
                                else:
                                    print("Too many high scoring matches!Ignoring!")
                                    print("Utterance: ", message["text"])
                                
                            elif num_high_matches == 0:
                                print("Warning! No match for surface form ", surface_form)
                                print("Utterance: ", message["text"])

                            else:
                                best_match = top_actor_matches[0][0]
                                cached_result = match_cache.get(best_match)

                                if cached_result:
                                    imdb_records.append(cached_result)
                                else:
                                    imdb_person_results = ia.search_person(best_match)
                                    if len(imdb_person_results) > 0:
                                        imdb_records.append(imdb_person_results[0])
                                        match_cache[best_match] = imdb_person_results[0]

            message["imdb_entries"] = [item.asXML() for item in imdb_records]
    return conversations

def load_movie_mentions_csv(mentions_csv_path):
    with open(mentions_csv_path, 'r') as csv_file:
        return [row for row in csv.DictReader(csv_file)]

def fix_movie_mentions(conversations, movie_corrected_mentions, imdb_sqlite_path):
    if imdb_sqlite_path:
        ia = imdb.IMDb('s3', os.path.join('sqlite+pysqlite:///', args.imdb_sqlite_path))
    else:
        ia = imdb.IMDb()

    def extract_imdb_id(url: str):
        return url.replace("https://www.imdb.com/title/tt", "").replace("/", "")

    for mention in movie_corrected_mentions:
        conv_id = mention["conversation_index"]
        message_index = mention["message_index"]

        conversation = conversations[int(conv_id)]

        message = conversation["messages"][int(message_index)]

        message_text = message["text"]

        movie_mentions_dict = conversation["movieMentions"]



        if mention["Movie Name 1"]:
            mention1 = mention["Movie Name 1"]
            imdb_id_1 = extract_imdb_id(mention["IMDB ID 1"])
            message_text = message_text.replace(mention1, imdb_id_1)

            movie_1 = ia.get_movie(imdb_id_1)
            movie_mentions_dict[imdb_id_1] = movie_1['title']


        if mention["Movie Name 2"]:
            mention2 = mention["Movie Name 2"]
            imdb_id_2 = extract_imdb_id(mention["IMDB ID 2"])

            movie_2 = ia.get_movie(imdb_id_2)
            movie_mentions_dict[imdb_id_2] = movie_2['title']

            message_text = message_text.replace(mention2, imdb_id_2)

        if mention["Movie Name 3"]:
            mention3 = mention["Movie Name 3"]
            imdb_id_3 = extract_imdb_id(mention["IMDB ID 3"])

            movie_3 = ia.get_movie(imdb_id_3)
            movie_mentions_dict[imdb_id_3] = movie_3['title']
            message_text = message_text.replace(mention3, imdb_id_3)

        message["text"] = message_text


        # Special corner case for correcting data
        if conv_id == "504" and message_index == "11" :
            message["text"] = message_text.replace("@Adam", "Adam")

    return conversations



def add_genre_mentions(conversations):
    ner_tagger = load_ner_tagger("gez/genre-phrases.tsv", "gez/movie-lex.json")

    for conversation in tqdm(conversations):
        for message in conversation["messages"]:
            tokenized_text = ner_tagger.tokenize_text(message["text"])

            mentions = ner_tagger.tag_tokens(tokenized_text)

            message["genre_mentions"] = mentions


    return conversations


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dataset_path', 
        type=str,
        default="redial"
    )
    parser.add_argument(
        '--imdb_sqlite_path',
        type=str,
        default=None,
        help="Path to the IMDB sqlite database. Optional, but makes search significantly faster"
    )

    args = parser.parse_args()

    splits = {
        "train": {"input_file": "train_data_imdb_corrected.jsonl", "output_file": "train_data_genre_tagged.jsonl"},
        "test": {"input_file": "test_data_imdb_corrected.jsonl", "output_file": "test_data_genre_tagged.jsonl"}
    }
    
    for split, metadata in splits.items():

        split_filepath = os.path.join(args.dataset_path, metadata["input_file"])

        conversations = load_conversations(split_filepath)

        annotated_conversations = add_genre_mentions(conversations)

        out_filepath = os.path.join(args.dataset_path, metadata["output_file"])
        dump_conversations_to_file(annotated_conversations, out_filepath)<file_sep>"""
 Created by diesel
 11/9/20
"""


import pandas as pd
import json


d = {
    'movieMentions': {
        '111776': 'Super Troopers (2001)',
        '91481': '<NAME> (1984)',
        '151656': 'Police Academy  (1984)',
        '134643': 'American Pie  (1999)',
        '192131': 'American Pie ',
        '124771': '48 Hrs. (1982)',
        '94688': 'Police Academy 2: Their First Assignment (1985)',
        '101794': 'Lethal Weapon (1987)'
    },
    'respondentQuestions': {
        '111776': {
            'suggested': 0,
            'seen': 1,
            'liked': 1
        },
        '91481': {
            'suggested': 1,
            'seen': 2,
            'liked': 2
        },
        '151656': {
            'suggested': 1,
            'seen': 0,
            'liked': 1
        },
        '134643': {
            'suggested': 0,
            'seen': 1,
            'liked': 1
        },
        '192131': {
            'suggested': 0,
            'seen': 1,
            'liked': 1
        },
        '124771': {
            'suggested': 1,
            'seen': 2,
            'liked': 2
        }, '94688': {
            'suggested': 1,
            'seen': 0,
            'liked': 1
        },
        '101794': {
            'suggested': 1,
            'seen': 0,
            'liked': 2
        }
    },
    'messages': [
        {
            'timeOffset': 0,
            'text': 'Hi I am looking for a movie like @111776',
             'senderWorkerId': 956,
            'messageId': 204171
        },
        {
            'timeOffset': 48,
            'text': 'You should watch @151656',
            'senderWorkerId': 957,
            'messageId': 204172
        },
        {
            'timeOffset': 90,
            'text': 'Is that a great one? I have never seen it. I have seen @192131',
            'senderWorkerId': 956,
            'messageId': 204173
        },
        {
            'timeOffset': 122,
            'text': 'I mean @134643',
            'senderWorkerId': 956,
            'messageId': 204174
        },
        {
            'timeOffset': 180,
            'text': 'Yes @151656 is very funny and so is @94688',
            'senderWorkerId': 957,
            'messageId': 204175
        },
        {
            'timeOffset': 199,
            'text': 'It sounds like I need to check them out',
            'senderWorkerId': 956,
            'messageId': 204176
        },
        {
            'timeOffset': 219,
            'text': 'yes you will enjoy them',
            'senderWorkerId': 957,
            'messageId': 204177
        },
        {
            'timeOffset': 253,
            'text': 'I appreciate your time. I will need to check those out. Are there any others you would recommend?',
            'senderWorkerId': 956,
            'messageId': 204178
        },
        {
            'timeOffset': 297,
            'text': 'yes @101794',
             'senderWorkerId': 957,
            'messageId': 204179
        },
        {
            'timeOffset': 311,
            'text': 'Thank you i will watch that too',
             'senderWorkerId': 956,
            'messageId': 204180
        },
        {
            'timeOffset': 312,
            'text': 'and also @91481',
             'senderWorkerId': 957,
            'messageId': 204181
        },
        {
            'timeOffset': 326,
            'text': 'Thanks for the suggestions.',
             'senderWorkerId': 956,
            'messageId': 204182
        },
        {
            'timeOffset': 341,
            'text': 'you are welcome',
             'senderWorkerId': 957,
            'messageId': 204183
        },
        {
            'timeOffset': 408,
            'text': 'and also @124771',
             'senderWorkerId': 957,
            'messageId': 204184
        },
        {
            'timeOffset': 518,
            'text': 'thanks goodbye',
             'senderWorkerId': 956,
            'messageId': 204185
        }
    ],
    'conversationId': '20001',
    'respondentWorkerId': 957,
    'initiatorWorkerId': 956,
    'initiatorQuestions': {
        '111776': {
            'suggested': 0, 'seen': 1, 'liked': 1},
        '91481': {
            'suggested': 1, 'seen': 2, 'liked': 2},
        '151656': {
            'suggested': 1, 'seen': 0, 'liked': 1},
        '134643': {
            'suggested': 0, 'seen': 1, 'liked': 1},
        '192131': {
            'suggested': 0, 'seen': 1, 'liked': 1},
        '124771': {
            'suggested': 1, 'seen': 2, 'liked': 2},
        '94688': {
            'suggested': 1, 'seen': 0, 'liked': 1},
        '101794': {
            'suggested': 0, 'seen': 2, 'liked': 2}}}



def get_messages(infile):
    with open(infile, "r") as fin:
        messages = []
        for line in fin:
            d = json.loads(line)

            speaker_lookup = dict(zip(range(len("ABCDEF")), "ABCDEF"))
            speakers = {}
            sid = 0
            for m in d["messages"]:
                #assert m['senderWorkerId'] < 7
                if m['senderWorkerId'] not in speakers:
                    speakers[m['senderWorkerId']] = speaker_lookup[sid]
                    sid += 1

                messages.append({
                    "text": m["text"],
                    "speaker": speakers[m['senderWorkerId']]
                })
    return messages



class DataLoader(object):
    def __init__(self, args, infile=None):
        self.args = args
        self.infile = infile



    def load(self, infile=None):
        if infile is None:
            infile = self.infile
        else:
            self.infile = infile

        messages = get_messages(infile)
        return pd.DataFrame(messages)



from lexicon import Lexicon, LexBuilder
from rule_based_ner import RuleBasedNER
import file_utils as fu
import os
from collections import defaultdict

def load_ner_tagger(concept_path, lex_save_path):
    df = pd.read_csv(concept_path, sep="\t")
    builder = LexBuilder()

    key_phrases = defaultdict(list)
    for row in df.itertuples():
        key_phrases[row.group.strip()].append(row.phrase.strip())

    for standard_form, forms in key_phrases.items():
        # entry is dictionary
        builder.add_entry({
            "category": "movie_genre",
            "standard_form": standard_form,
            "forms": forms,
            "name": standard_form,
            "full_name": standard_form,
        })

    builder.build_lexes()
    #builder.save_lexes(os.path.join("./", f"{topic_name}-lex.json"))
    builder.save_lexes(lex_save_path)

    #print("entry keys:", builder.get_entry_key_set())

    return RuleBasedNER(builder.lexicons)





def main():

    infile = "redial/train_data.jsonl"

    loader = DataLoader(None)
    messages = loader.load(infile)

    ner_tagger = load_ner_tagger("gez/genre-phrases.tsv", "gez/movie-lex.json")

    mentions = []
    for row in messages.itertuples():
        #print("\ntext:", row.text)
        toks = ner_tagger.tokenize_text(row.text)
        mentions.append(json.dumps(ner_tagger.tag_tokens(toks)))

    messages["genre_mentions"] = mentions

    #df = pd.DataFrame(messages)
    #df.to_csv("redial/train_data.messages.tsv")
    messages.to_csv("redial/train-genre-mentions.tsv", sep="\t", index=False)


if __name__ == "__main__":
    main()
<file_sep>import pickle
from imdb import IMDb
import pandas as pd
from copy import deepcopy
import json

from imdb.Person import Person
from imdb.Movie import Movie
from imdb.Company import Company

from tqdm import tqdm


ia = IMDb()



def dump_person(p, max_n=None):
    """
    print("person:", p)
    print("current_info:", p.current_info)
    print("infoset2keys:", p.infoset2keys)
    d = dict(p)

    print(d)
    print("\n\n")
    for name, attr_set in p.infoset2keys.items():
        for attr in attr_set:
            print(f"{attr} = {p.get(attr)}")

    print("\n\n", "-"*20)
    #p = ia.get
    ia.update(p, info=["main", "biography", "awards"])

    print("person:", p)
    print("current_info:", p.current_info)
    print("infoset2keys:", p.infoset2keys)
    d = dict(p)

    print(d)
    print("\n\n")
    for name, attr_set in p.infoset2keys.items():
        print(name.upper())
        for attr in attr_set:
            r = p.get(attr)
            if isinstance(r, dict):
                r = r.keys()
            print(f"{attr} = {r}")
    """
    print("  * _person_", p)
    if p.myID not in all_people:
        ia.update(p, info=["main", "biography", "awards"])
        all_people[p.myID] = dump_item(dict(p), max_n)

    return p.myID

def dump_company(c, max_n=None):
    print("  * _company_", c)
    if c.myID not in all_companies:
        ia.update(c, info=["main",])
        all_companies[c.myID] = dump_item(dict(c), max_n)
    return c.myID

def dump_movie(m, max_n=None):
    print("  * _movie_", m)
    if m.movieID not in all_movies:
        all_movies[m.movieID] = True
        temp_d = dict(m)
        movie_d = {}
        for k, attr in temp_d.items():
            movie_d[k] = dump_item(attr, max_n)
        all_movies[m.movieID] = movie_d

    return m.movieID

def dump_item(val, max_n=None):

    if isinstance(val, list):
        trunc_val = len(val) if max_n is None else max_n
        dumped = [dump_item(v, max_n) for v in val[:trunc_val]]
    elif isinstance(val, dict):
        dumped = {k: dump_item(v, max_n) for k,v in val.items()}
    elif isinstance(val, Person):
        dumped = dump_person(val, max_n)
    elif isinstance(val, Movie):
        dumped = dump_movie(val, max_n)
    elif isinstance(val, Company):
        dumped = dump_company(val, max_n)
    else:
        # base case
        dumped = val
    return dumped





def main():
    print("loading data ...")
    with open("imdb_data.pkl", "rb") as fin:
        imdb_data = pickle.load(fin)

    print("aggregating data ...")
    for k, movie in tqdm(list(imdb_data["with_imdb_key"].items())):
        all_movies[k] = dump_item(movie, 1)
        break

    print("saving data ...")
    for name, data in {"movies": all_movies, "people": all_people, "companies": all_companies}.items():
        with open(f"all-{name}.json", "w") as fout:
            json.dump(data, fout, indent=2)


def demo_main():
    #ia = IMDb()
    print("ia.get_movie_infoset():", ia.get_movie_infoset())
    print("ia.get_person_infoset():", ia.get_person_infoset())
    print("ia.get_company_infoset():", ia.get_company_infoset())
    exit(0)
    with open("imdb_data.pkl", "rb") as fin:
        d = pickle.load(fin)
    all_movies = {}
    all_people = {}
    for k, v in d["with_imdb_key"].items():
        # print(f"{k}: {type(v)}")
        print(f"{v.myTitle}:", v.current_info)
        print(v.infoset2keys)
        print(v.get("title"))
        print(v.get("taglines"))
        print("median" in v)

        ia.update(v, info=['taglines', 'vote details'])
        print(f"{v.myTitle}:", v.current_info)
        print(v.infoset2keys)
        print(v.get("title"))
        print(v.get("taglines"))
        print("median" in v)
        for a in ['demographics', 'number of votes', 'arithmetic mean', 'median']:
            print(f"{a}:", v.get(a))

        for kk, vv in v.items():
            # print(f"{kk}: {vv}")
            # if kk == "cast":
            if False:
                p = vv[0]
                print("-" * 10)
                print("cast[0]")
                print(p.summary())
                print(p.myName)
                print(p.personID)
                print(p.billingPos)
                for a in p.__dict__.keys():
                    print(f"{a}: {p.get(a)}")
                print("-" * 10)

                for attr_name in p.keys_alias:
                    print(f"{attr_name}: {p.get(attr_name)}")
                print("-" * 20)

                p = dict(name='', personID=p.personID, myName=p.myName,
                         myID=p.myID, data=p.data,
                         currentRole=p.currentRole,
                         roleIsPerson=p._roleIsPerson,
                         notes=p.notes, accessSystem=p.accessSystem,
                         titlesRefs=p.titlesRefs,
                         namesRefs=p.namesRefs,
                         charactersRefs=p.charactersRefs)
                for _k, _v in p.items():
                    print(f"{_k}: {_v}")

        _d = dict(v)
        print(_d)
        if input(">"):
            exit(0)
    df = pd.DataFrame([dict(movie) for movie in d["with_imdb_key"]])
    df.to_csv("imdb-movies.tsv", sep="\t", index=False)
    print("df shape:", df.shape)

if __name__ == "__main__":
    #demo_main()
    main()

<file_sep>from .redial import RedialDialoGPTDataset, RedialTransferTransfoDataset<file_sep>
"""
This script allows for the standardized evaluation metrics
such as F1, BLEU-4, ROUGE-L, METEOR, etc. to be computed
against a set of reference responses.
"""
import argparse
import glob
import logging
import os
import pickle
import pprint
from collections import Counter

import nlgeval
from datasets import load_metric
import nltk
from nltk.translate.meteor_score import meteor_score
import numpy as np

class ReferenceMetric(object):
    """
    Metric that requires a reference sentence for each
    hypothesis to compute
    """
    def compute(self, hypotheses, references):
        raise NotImplementedError("Implement the compute method!")


class ReferenceFreeMetric(object):
    """
    Metric that does not require a reference sentence
    """

    def compute(self, hypotheses):
        raise NotImplementedError("Implement the compute method!")


class NLPReferenceMetric(ReferenceMetric):
    """
    Reference dependent metrics that are part of the
    Huggingface NLP library
    """
    def __init__(self, module, compute_args={}):
        self.scorer = load_metric(module)
        self.compute_args = compute_args

    def compute(self, hypotheses, references):
        return self.scorer.compute(predictions=hypotheses, references=references, **self.compute_args)


class BLEUMetric(NLPReferenceMetric):
    def __init__(self):
        super().__init__('bleu')

    def __repr__(self):
        return 'BLEU-4'

    def compute(self, hypotheses, references):
        return super().compute([hyp.split() for hyp in hypotheses], [[ref.split()] for ref in references])


class RougeMetric(NLPReferenceMetric):
    def __init__(self):
        super().__init__('rouge')

    def __repr__(self):
        return 'ROUGE'

class BertScoreMetric(NLPReferenceMetric):
    def __init__(self):
        self.arg_dict = {"lang": "en"}
        super().__init__('bertscore', self.arg_dict)

    def __repr__(self):
        return f'BertScore({self.arg_dict})'

    def compute(self, hypotheses, references):
        return sum(super().compute(hypotheses, references)['f1']) / len(hypotheses)


class UnigramFScoreMetric(ReferenceMetric):
    def __init__(self, beta=1):
        self.beta = beta
        self.beta_squared = beta * beta  # Fbeta uses (beta^2 as weighting factor)

    def _f_beta(self, pred, true):
        common = Counter(true) & Counter(pred)
        num_same = sum(common.values())

        if num_same == 0:
            return 0

        prec = num_same / len(pred)
        rec = num_same / len(true)

        f_beta = ((self.beta_squared + 1) * prec * rec) / (self.beta_squared * prec + rec)
        return f_beta

    def compute(self, hypotheses, references):
        return sum([self._f_beta(hyp.split(), ref.split()) for hyp, ref in zip(hypotheses, references)]) / len(references)

    def __repr__(self):
        return f'F{self.beta}-score'

class MeteorMetric(ReferenceMetric):
    """
    Computes the average METEOR score across all examples
    """

    def compute(self, hypotheses, references):

        try:
            nltk.data.find('wordnet')
        except LookupError:
            nltk.download('wordnet')

        return sum([meteor_score([ref], hyp) for (ref, hyp) in zip(hypotheses, references)]) / len(references)

    def __repr__(self):
        return f'METEOR score'

class NGramDiversity(ReferenceFreeMetric):
    def __init__(self, n=1):
        self.n = n

    def _diversity(self, pred):
        n_grams = []

        for i in range(len(pred) - self.n + 1):
            n_grams.append(' '.join(pred[i:i + self.n]))

        if len(n_grams) == 0:
            return 0

        return len(set(n_grams)) / len(n_grams)

    def compute(self, hypotheses):
        return sum([self._diversity(hyp.split()) for hyp in hypotheses]) / len(hypotheses)

    def __repr__(self):
        return f'Utterance {self.n}-gram diversity'

class CorpusNGramDiversity(ReferenceFreeMetric):
    """
    Computes the number of unique n-grams that were generated by the model
    out of the total number of n-grams present
    """
    def __init__(self, n=1):
        self.n = n

    def compute(self, hypotheses):
        n_grams = []

        for hyp in hypotheses:
            pred = hyp.split()
            for i in range(len(pred) - self.n + 1):
                n_grams.append(' '.join(pred[i:i + self.n]))

        if len(n_grams) == 0:
            return 0

        return len(set(n_grams)) / len(n_grams)

    def __repr__(self):
        return f'Corpus {self.n}-gram diversity'

class NLGEval(ReferenceMetric):
    """
    Runs the full NLGEval pipeline which computes multiple machine translation
    metrics:
    1. BLEU
    2. METEOR
    3. ROUGE
    3. CIDEr
    4. Skip Thought
    5. Embedding Average
    6. Vector Extrema
    7. Greedy matching
    """

    def __init__(self, hypothesis_file, reference_file):
        self.metrics_dict = nlgeval.compute_metrics(hypothesis=hypothesis_file,
                                                    references=[reference_file], no_skipthoughts=True, no_glove=True)
    def __repr__(self):
        return 'NLGEval metrics'

    def compute(self, hypotheses, references):
        return pprint.pformat(self.metrics_dict)

class USRMetric(ReferenceFreeMetric):
    """
    Computes the USR score as defined by Mehri and Eskenazi (2019)

    Currently the code computes the following:
    1. MLM objective
    """
    def _make_scoring_file(self, context_file, hypothesis_file):

        scratch_path = '../submissions/scratch'
        with open(context_file, 'r') as context_, open(hypothesis_file, 'r') as hypothesis_, open(scratch_path, 'w') as scratch_file:
            context_lines = [line.strip() for line in context_]
            hypothesis_lines = [line.strip() for line in hypothesis_]

            for (c, h) in zip(context_lines, hypothesis_lines):
                scratch_file.writelines(f'{c} _eos _go {h}\n')

        return scratch_path

    def _compute_regression_scores(self, mlm_score, dr_c_scores, dr_f_scores):
        # Understandable (MLM), Natural (MLM), Maintains Context (DR-c), Interesting (DR-c), Uses Knowledge (DR-f)
        common_len = min(len(mlm_score), len(dr_c_scores), len(dr_f_scores))

        mlm_score = mlm_score[:common_len]
        dr_c_scores = dr_c_scores[:common_len].tolist()
        dr_f_scores = dr_f_scores[:common_len].tolist()

        X = np.array([mlm_score, mlm_score, dr_c_scores, dr_c_scores, dr_f_scores]).T

        with open('../usr/examples/regr.pkl', 'rb') as regression_model_file:
            model = pickle.load(regression_model_file)
        y = model.predict(X)

        return y.tolist()


    def _compute_dr_score(self, args):
        from usr.examples.train_understandable import MODEL_CLASSES, WEIGHTS_NAME, predict

        config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
        checkpoints = [args.output_dir]
        preds = None
        if args.eval_all_checkpoints:
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))
            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        for checkpoint in checkpoints:
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
            prefix = checkpoint.split('/')[-1] if checkpoint.find('checkpoint') != -1 else ""

            model = model_class.from_pretrained(checkpoint)
            model.to(args.device)

            preds = predict(args, model, tokenizer, args.data_dir, prefix=prefix)

        return preds

    def make_args(self, model_path, scoring_file):
        args = argparse.Namespace()
        args.per_gpu_eval_batch_size = 1
        args.output_dir = model_path
        args.model_type = "roberta"
        args.model_name_or_path = "roberta-base"
        args.data_dir = scoring_file
        args.do_eval = True
        args.task_name = "qqp"
        args.do_lower_case = False
        args.device = "cuda"
        args.eval_all_checkpoints = False
        args.max_seq_length = 128
        args.local_rank = -1
        args.n_gpu = 1
        args.overwrite_cache = True
        args.output_mode = 'classification' # QQP task uses classification
        return args

    def _compute_dr_c_score(self, scoring_file):
        args = self.make_args("../usr/examples/ctx", scoring_file)
        preds = self._compute_dr_score(args)
        return preds

    def _compute_dr_f_score(self, scoring_file):
        args = self.make_args("../usr/examples/uk", scoring_file)
        return self._compute_dr_score(args)

    def __init__(self, context_file, fact_file, hypothesis_file):
        self.lm_scoring_file = self._make_scoring_file(context_file, hypothesis_file)
        self.dr_scoring_file = self._make_dr_scoring_file(context_file, fact_file, hypothesis_file)


    def compute_mlm_scores(self, scoring_file):
        args = self.build_args(scoring_file)
        from usr.examples.run_lm_finetuning import evaluate, MODEL_CLASSES, WEIGHTS_NAME
        config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
        config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
                                              cache_dir=args.cache_dir if args.cache_dir else None)
        tokenizer = tokenizer_class.from_pretrained(
            args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
            do_lower_case=args.do_lower_case,
            cache_dir=args.cache_dir if args.cache_dir else None)
        if args.block_size <= 0:
            args.block_size = tokenizer.max_len_single_sentence
        model = model_class.from_pretrained(args.model_name_or_path,
                                            from_tf=bool('.ckpt' in args.model_name_or_path),
                                            config=config,
                                            cache_dir=args.cache_dir if args.cache_dir else None)
        model.to(args.device)
        result = None
        checkpoints = [args.output_dir]
        if args.eval_all_checkpoints:
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))
            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        for checkpoint in checkpoints:
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
            prefix = checkpoint.split('/')[-1] if checkpoint.find('checkpoint') != -1 else ""

            model = model_class.from_pretrained(checkpoint)
            model.to(args.device)
            result = evaluate(args, model, tokenizer, prefix=prefix)

            # result = dict((k + '_{}'.format(global_step), v) for k, v in result.items())
            # results.update(result)
        return result

    def build_args(self, scoring_file):
        args = argparse.Namespace()
        args.output_dir = '../usr/examples/roberta_ft'
        args.model_type = 'roberta'
        args.train_data_file = scoring_file
        args.per_gpu_eval_batch_size = 1
        args.model_name_or_path = 'roberta-base'
        args.eval_data_file = scoring_file
        args.do_eval = True
        args.mlm = True
        args.device = "cuda"
        args.local_rank = -1
        args.n_gpu = 0
        args.config_name = ""
        args.cache_dir = ""
        args.tokenizer_name = ""
        args.do_lower_case = False
        args.eval_all_checkpoints = False
        args.block_size = -1
        return args

    def __repr__(self):
        return 'USR Metric Fine-tuned on Topical Chats'

    def compute(self, hypotheses):
        mlm_scores = self.compute_mlm_scores(self.lm_scoring_file)
        dr_c_scores = self._compute_dr_c_score(self.dr_scoring_file)
        dr_f_scores = self._compute_dr_f_score(self.dr_scoring_file)
        usr_scores = self._compute_regression_scores(mlm_scores, dr_c_scores, dr_f_scores)

        self.results = usr_scores

        return sum(self.results) / len(self.results) if len(self.results) > 0 else 0

    def _make_dr_scoring_file(self, context_file, fact_file, hypothesis_file):
        scratch_dir = '../submissions/'
        scratch_path = scratch_dir + 'dev.tsv'
        with open(context_file, 'r') as ctx, open(fact_file, 'r') as fct, open(hypothesis_file) as resp, open(scratch_path, 'w') as scratch_file:
            contexts = [line.strip() for line in ctx]
            facts = [line.strip() for line in fct]
            responses = [line.strip() for line in resp]

            for (c, f, r) in zip(contexts, facts, responses):
                scratch_file.write(f'0\t1\t2\t{c} {f} _eos\t_go {r}\t0\n')
        return scratch_dir

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # Currently support only single hypotheses scoring
    parser.add_argument('--context_file',
                        type=str,
                        default='processed_output/valid_freq.src')

    parser.add_argument('--predictions_file',
                        type=str,
                        default="submissions/submissions.txt",
                        help='File containing output predictions')
    parser.add_argument('--references_file',
                        type=str,
                        default='processed_output/valid_freq.tgt',
                        help='File containing the reference responses')
    parser.add_argument('--fact_file',
                        type=str,
                        default='processed_output/valid_freq.fct',
                        help='File containing reference facts')

    args = parser.parse_args()

    with open(args.predictions_file, 'r') as predictions_file:
        predictions = [line
                           .strip()
                           .replace(".", " .").replace("?", " ?")
                           .replace(",", " ,")
                           .replace("'", " ' ")
                           .replace("dn't", "d n't")
                       for line in predictions_file]

    with open(args.references_file, 'r') as references_file:
        references = [line.replace("_go", "").replace("_eos", "").strip() for line in references_file]

    assert len(predictions) == len(references), "The number of predictions and references do not match!"

    for prediction in predictions:
        assert prediction.strip() != "", "Predictions cannot be empty!"

    metrics = [
        # BLEUMetric(),
        # RougeMetric(),
        BertScoreMetric(),
        # MeteorMetric(),
        UnigramFScoreMetric(),
        NGramDiversity(n=1),
        NGramDiversity(n=2),
        CorpusNGramDiversity(n=1),
        CorpusNGramDiversity(n=2),
        NLGEval(args.predictions_file, args.references_file),
        # USRMetric(args.context_file, args.fact_file, args.predictions_file)
    ]

    print(f"Number of examples n={len(predictions)}\n")

    for metric in metrics:
        if isinstance(metric, ReferenceFreeMetric):
            print(metric, ":")
            print(metric.compute(predictions))
        else:
            print(metric, ":")
            print(metric.compute(predictions, references))
<file_sep>import json
import argparse
import re
import csv
import os

from tqdm.auto import tqdm

from data_utils import load_conversations

def identify_mislabeled_movies(conversations, output_file_path):
    mislabeled_movie_regex = re.compile(r"@[A-Za-z]")

    mislabeled_movie_utterances = []


    for i, conversation in enumerate(tqdm(conversations)):
        for j, message in enumerate(conversation["messages"]):
            
            match = mislabeled_movie_regex.search(message["text"], re.MULTILINE)

            if match:
                mislabeled_movie_utterances.append({
                    "conversation_index": i,
                    "message_index": j,
                    "messageId": message["messageId"],
                    "text": message["text"]
                })

    
    with open(output_file_path, 'w') as mislabeled_utterances_file:
        writer = csv.DictWriter(mislabeled_utterances_file, fieldnames=["conversation_index", "message_index", "text", "messageId"])

        writer.writeheader()
        writer.writerows(mislabeled_movie_utterances)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('data_path')

    args = parser.parse_args()

    splits = {
        "train": {"input_file": "train_data_spacy.jsonl", "output_file": "train_utterances_with_mislabeled_movies.csv"},
        "test": {"input_file": "test_data_spacy.jsonl", "output_file": "test_utterances_with_mislabeled_movies.csv"}
    }

    for splitname, split_metadata in splits.items():

        input_file_path = os.path.join(args.data_path, split_metadata["input_file"])
        output_file_path = os.path.join(args.data_path, split_metadata["output_file"])
        conversations = load_conversations(input_file_path)

        identify_mislabeled_movies(conversations, output_file_path)<file_sep>import numbers
from collections import Sequence

from torch.optim.lr_scheduler import _LRScheduler

class PiecewiseLinearLR(_LRScheduler):
    """
    LR Scheduler that linearly decreases the LR over the
    duration of the number of steps, based off the Ignite
    implementation.

    I didn't want to change any of the original setup from TransferTransfo
    therefore, it felt more appropriate to retrofit a custom scheduler
    """

    def __init__(self, optimizer, milestones_values, last_epoch=-1):

        values = []
        milestones = []

        for pair in milestones_values:
            if not isinstance(pair, Sequence) or len(pair) != 2:
                raise ValueError("Argument milestones_values should be a list of pairs (milestone, param_value)")
            if not isinstance(pair[0], numbers.Integral):
                raise ValueError("Value of a milestone should be integer, but given {}".format(type(pair[0])))
            if len(milestones) > 0 and pair[0] < milestones[-1]:
                raise ValueError("Milestones should be increasing integers, but given {} is smaller "
                                 "than the previous milestone {}".format(pair[0], milestones[-1]))
            milestones.append(pair[0])
            values.append(pair[1])

        self.values = values
        self.milestones = milestones
        self._index = 0
        self.last_epoch = last_epoch
        super(PiecewiseLinearLR, self).__init__(optimizer, last_epoch)

    def _get_start_end(self):
        if self.milestones[0] > self.last_epoch:
            return self.last_epoch - 1, self.last_epoch, self.values[0], self.values[0]
        elif self.milestones[-1] <= self.last_epoch:
            return self.last_epoch, self.last_epoch + 1, self.values[-1], self.values[-1],
        elif self.milestones[self._index] <= self.last_epoch < self.milestones[self._index + 1]:
            return self.milestones[self._index], self.milestones[self._index + 1], \
                   self.values[self._index], self.values[self._index + 1]
        else:
            self._index += 1
            return self._get_start_end()

    def get_lr(self):
        start_index, end_index, start_value, end_value = self._get_start_end()
        return [start_value + (end_value - start_value) * (self.last_epoch - start_index) / (end_index - start_index)]
<file_sep>class GlobalStepCounter(object):
    def __init__(self):
        self.num_steps = 0

    def get(self):
        return self.num_steps

    def step(self):
        self.num_steps += 1<file_sep>"""
 Created by diesel
 11/11/20
"""


def main():
    with open("genres.txt", "r") as fin:
        lines = [l for l in fin.read().split("\n") if l]

    all_terms = []
    for line in lines:
        phrases = line.split(";")
        all_terms.extend([p.strip() for p in phrases if p.strip()])

    all_terms = list(set(all_terms))

    with open("genres-clean.txt", "w") as fout:
        fout.write("\n".join(all_terms))

if __name__ == "__main__":
    main()
<file_sep>"""
 Created by diesel
 11/9/20
"""





from lexicon import Lexicon, LexBuilder
from rule_based_ner import RuleBasedNER


import file_utils as fu
from shutil import copyfile

import unidecode

import os
import json
import nltk
from collections import Counter
from collections import defaultdict
import pandas as pd

import spacy
nlp = spacy.load('en_core_web_sm')



def unicode_2_ascii(text):
    #if isinstance(text, unicode):
    text = unidecode.unidecode(text)
    return text

def iter_lines(lines):
    for line in lines:
        yield line



def unpack_entry(first, lines, val_lists=False):
    # first line is the etype
    #if not first:

    d = {
        "category": first.strip(" :").lower()
    }
    single_string_values = {"text", "follow_up", "followup", "follow_up", "tid", "evaluative_language",
                            "dact", "etype", "mod", "rel", "primary_dact", "uid"}


    #print("first:", first)
    #print("lines[0]:", lines[0])
    for line in lines:
        if line.startswith("#"):
            continue
        line = line.strip()
        # stop at first empty line
        if not line:
            break

        splitline = line.split(":", maxsplit=1)
        #print("splitline:", splitline)

        slot_name, slot_vals = splitline
        slot_vals = unicode_2_ascii(slot_vals)

        if (val_lists and slot_name not in single_string_values) or slot_name in {"reference_forms", "forms", "referential_expressions"}:
            if "," in slot_vals:
                slot_vals = [v.strip() for v in slot_vals.split(",") if v.strip()]
                #print("ifif:", slot_vals)
            else:
                slot_vals = [v.strip() for v in slot_vals.split(";") if v.strip()]
                #print("ifelse")
        else:
            slot_vals = slot_vals.strip()
            #print("else")
        slot_name = slot_name.strip().lower()
        if slot_name == "tid":
            slot_name = "uid"
        elif slot_name == "followup":
            slot_name = "follow_up"
        d[slot_name] = slot_vals
        #print(f"{slot_name}: {slot_vals}")

    return d, d["category"]


def extract_referencial_forms(e, cat, fullname_only=False):
    if cat == "person":
        names = [e.get(name) for name in ["first_name", "last_name", ] if e.get(name)]
        full_name = " ".join(names)
        forms = [full_name]
        if not fullname_only and e.get("last_name"):
            forms.append(e["last_name"])
        e["full_name"] = full_name

    elif cat == "book":
        forms = [e[attr] for attr in ["full_name", "nick_name" ] if e.get(attr)]
        if e.get("nick_name") and e["nick_name"].startswith("the "):
            forms.append(e["nick_name"][len("the "):].strip())
        _ref_field = "reference_forms"
        if e.get(_ref_field):
            if not isinstance(e[_ref_field], list):
                e[_ref_field] = [e[_ref_field]]
            forms.extend(e[_ref_field])
    else:
        forms = [e[name] for name in ["name", "nick_name", "full_name",] if e.get(name)]
        if e.get("forms"):
            forms.extend(e["forms"])
        forms = list(set(forms))
        if not e.get("full_name") and e.get("name"):
            e["full_name"] = e["name"]

    #print("e forms:", e["forms"])
    if forms:
        e["forms"] = forms
    return e


def demo():
    center_dir = "../center/astext-cleaned/"
    topic_name = "animals"
    #fact_infile_name = f"{topic_name}-facts.txt"
    lex_infile_name = f"{topic_name}-lex.txt"
    lex_input_path = os.path.join(center_dir, lex_infile_name)
    lines = fu.read_lines(lex_input_path)

    lines = iter_lines(lines)
    builder = LexBuilder()



    # entry is dictionary
    builder.add_entry({
        'category': 'animal',
        'standard_form': 'Cat',
        'forms': ['kittens', 'kitten', 'cats', 'feline', 'felines', 'kitty', 'kitty cats', 'kitty cat', 'Cat', 'kitties'],
        'name': 'Cat',
        'full_name': 'Cat'
    })
    builder.add_entry({
        'category': 'animal',
        'standard_form': 'Dog',
        'forms': ['Dog', 'canine', 'doggy', 'dogs', 'puppies', 'puppy'],
        'name': 'Dog',
        'full_name': 'Dog'
    })
    builder.add_entry({
        'category': 'concept',
        'standard_form': 'gender',
        'forms': ['gender'],
        'name': 'gender',
        'full_name': 'gender'
    })
    builder.add_entry({
        'category': 'concept',
        'standard_form': 'intelligence',
        'forms': ['learning', 'intelligence', 'cognitive ability'],
        'name': 'intelligence',
        'full_name': 'intelligence'
    })
    builder.add_entry({
        'category': 'concept',
        'standard_form': 'anatomy',
        'forms': ['anatomy'],
        'name': 'anatomy',
        'full_name': 'anatomy'
    })

    builder.build_lexes()
    builder.save_lexes(os.path.join("./", f"{topic_name}-lex.json"))

    print("entry keys:", builder.get_entry_key_set())

    ner_tagger = RuleBasedNER(builder.lexicons)

    texts = ["i think cats dog", "i think cats learning gender and anatomy", "i think cats learning",
             "this is cats weird and dogs"]
    for text in texts:
        print("\ntext:", text)
        toks = ner_tagger.tokenize_text(text)
        mentions = ner_tagger.tag_tokens(toks)
        print("mentions:")
        for m in mentions:
            print(" - ", m)

        print()




def main():
    center_dir = "../center/astext-cleaned/"
    topic_name = "hockey"
    topic_name = "animals"
    #topic_name = "activities"
    #topic_name = "harry_potter"
    #fact_infile_name = f"{topic_name}-facts.txt"
    lex_infile_name = f"{topic_name}-lex.txt"
    lex_input_path = os.path.join(center_dir, lex_infile_name)
    lines = fu.read_lines(lex_input_path)

    lines = iter_lines(lines)
    builder = LexBuilder()

    for j, line in enumerate(lines):
        line = line.strip()
        if not line:
            continue
        if line.startswith("#"):
            continue
        entry, etype = unpack_entry(line, lines)
        # entry is dictionary
        entry = extract_referencial_forms(entry, etype, fullname_only=True)
        builder.add_entry(entry, etype)


    builder.build_lexes()
    builder.save_lexes(os.path.join("./", f"{topic_name}-lex.json"))

    print("entry keys:", builder.get_entry_key_set())

    ner_tagger = RuleBasedNER(builder.lexicons)

    texts = ["i think cats dog", "i think cats learning gender and anatomy", "i think cats learning",
             "this is cats weird and dogs"]
    for text in texts:
        print("\ntext:", text)
        toks = ner_tagger.tokenize_text(text)
        mentions = ner_tagger.tag_tokens(toks)
        print("mentions:")
        for m in mentions:
            print(" - ", m)

        print()






if __name__ == "__main__":
    #main()
    demo()
<file_sep>import csv
import imdb
import pickle
import argparse
import os
import requests
from collections import defaultdict
from tqdm.auto import tqdm, trange
from bs4 import BeautifulSoup


def retrieve_movies_data_from_imdb(args):
    """
    Given the merged movies data, retrieve data for ones which have a matching IMDB id
    """

    if args.imdb_sqlite_path:
        ia = imdb.IMDb('s3', os.path.join('sqlite+pysqlite:///', args.imdb_sqlite_path))
    else:
        ia = imdb.IMDb()
    movies_data = defaultdict(dict)

    with open(args.merged_movie_data_path, 'r') as merged_movies_file:
        reader = csv.DictReader(merged_movies_file)
        for row in tqdm(reader):
            database_id = row['databaseId']
            imdb_id = row['imdbId']
            
            if database_id != '-1' and imdb_id != '-1': # We can definitively identify and retrieve these movies
                try:
                    movies_data['with_imdb_key'][database_id] = movie = ia.get_movie(imdb_id)
                    print("Processed movie:", movie)
                except imdb.IMDbError as e:
                    print("Exception", e)
                    movies_data['without_imdb_key'][database_id] = row
                    print("Skipped movie", row['movieName'])
            else: # Let's kick the can down the road by dealing with it later
                movies_data['without_imdb_key'][database_id] = row
                print("Skipped movie", row['movieName'])

    return movies_data

def scrape_imdb_list(url, n_pages):
    item_list = []

    index = 1

    for i in trange(1, n_pages + 1):
        formatted_url = url.format(str(i))

        response = requests.get(formatted_url)
        if response.status_code == 200:
            html = response.text

            soup = BeautifulSoup(html, 'html.parser')

            for header in soup.find_all('h3', "lister-item-header"):
                item_name = header.a.get_text().strip()

                item_struct = {
                    'id': index,
                    'name': item_name
                }

                item_list.append(item_struct)

                index += 1

    headers = ['id', 'name']

    return (item_list, headers)

def scrape_imdb_top_1000_actors():
    url = "https://www.imdb.com/list/ls058011111/?sort=list_order,asc&mode=detail&page={}"


    actor_list, headers = scrape_imdb_list(url, 10)

    with open('top_1000_actors.csv', 'w') as top_actors_file:
        writer = csv.DictWriter(top_actors_file, fieldnames=headers)

        writer.writeheader()
        writer.writerows(actor_list)
    

def scrape_imdb_top_250_directors():
    url = "https://www.imdb.com/list/ls008344500/"

    director_list, headers = scrape_imdb_list(url, 3)
    with open('top_250_directors.csv', 'w') as top_directors_file:
        writer = csv.DictWriter(top_directors_file, fieldnames=headers)

        writer.writeheader()
        writer.writerows(director_list)

def scrape_imdb_top_500_directors():
    url = "https://www.imdb.com/list/ls039888167/"

    director_list, headers = scrape_imdb_list(url, 5)
    with open('top_500_directors.csv', 'w') as top_directors_file:
        writer = csv.DictWriter(top_directors_file, fieldnames=headers)

        writer.writeheader()
        writer.writerows(director_list)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--merged_movie_data_path', 
        default='redial/movies_merged_with_imdb.csv',
        type=str,
        help='Path to the merged file of imdb info and movielens data'
    )
    parser.add_argument('--imdb_sqlite_path',
        default='',
        type=str,
        help='Path to the IMDB sqlite data (for offline retrieval)'
    )

    args = parser.parse_args()

    # movies_data = retrieve_movies_data_from_imdb(args)

    # with open('imdb_data.pkl', 'wb') as movie_imdb_pickle_file:
    #     pickle.dump(movies_data, movie_imdb_pickle_file)

    # scrape_imdb_top_1000_actors()
    scrape_imdb_top_500_directors()<file_sep>import csv
import json
import os
import pickle

import re

import pdb
from bs4 import BeautifulSoup
from collections import namedtuple
from transformers import GPT2Tokenizer

from tqdm.auto import tqdm

BaselineExample = namedtuple(
    'BaselineExample',
    ['context', 'response']
)

KnowledgeGroundedExample = namedtuple(
    'KnowledgeGroundedExample',
    ['context', 'response', 'knowledge']
)



def prepare_baseline_redial_split(
    split_path,
    tokenizer,
    movie_db_map
):

    with open(split_path, 'r') as split_file:
        split_conversations = split_file.read().splitlines()


    examples = []

    # Matching for movie mention ids: @1234
    movie_mention_pattern = re.compile(r"@(\d+)")
    # Pattern for mathching the year portion: (2007)
    movie_title_year_pattern = re.compile(r"\s+\(\d+\)")
    for conversation_str in tqdm(split_conversations):
        conversation = json.loads(conversation_str)

        context = []

        messages = conversation["messages"]

        response = ""

        for i, message in enumerate(messages):
            processed_text = message["text"]

            for mention in movie_mention_pattern.finditer(processed_text):
                movie_id = mention.group(1)

                # Remove year from title
                movie_title = movie_title_year_pattern.sub('', movie_db_map[movie_id])
                # for now, naively substitute movie title in message
                processed_text = processed_text.replace("@" + movie_id, movie_title)

            if i == len(messages) - 1 or \
                message["senderWorkerId"] != messages[i + 1]["senderWorkerId"]:
                response += processed_text
                encoded_response = tokenizer.encode(response)
                examples.append(BaselineExample(
                    context,
                    encoded_response
                ))

                context = context + [encoded_response]
                response = ""
            else:
                # We looked ahead and saw another follow-on response
                response += processed_text + " . "

    return examples

def prepare_redial_knowledge_grounded_split(
        split_path,
        movie_db_map,
        recommender_only=False,
        include_dact=True
):
    print("\nLoading data", split_path)

    with open(split_path, 'r') as split_file:
        split_conversations = split_file.read().splitlines()

    examples = []

    # Matching for movie mention ids: @1234
    movie_mention_pattern = re.compile(r"@(\d+)")
    # Pattern for mathching the year portion: (2007)
    movie_title_year_pattern = re.compile(r"\s+\(\d+\)")

    num_examples_using_knowledge = 0


    unk_terms = {
        "movie_genre": "<unk_genre>",
        "director": "<unk_director>",
        "cast": "<unk_cast>",
    }
    dact_set = set()

    for conversation_str in tqdm(split_conversations):
        conversation = json.loads(conversation_str)

        context = []

        messages = conversation["messages"]
        response = ""
        response_knowledge = []

        if recommender_only:
            sender_id = conversation["initiatorWorkerId"]

        for i, message in enumerate(messages):

            processed_text = message["text"]

            for mention in movie_mention_pattern.finditer(processed_text):
                movie_id = mention.group(1)

                movie_title = movie_db_map.get(movie_id)
                if not movie_title:
                    movie_title = conversation["movieMentions"][movie_id]
                    if isinstance(movie_title, dict):
                        movie_title = movie_title.get("title")

                movie_title = movie_title_year_pattern.sub('', movie_title)
                # naively substitute movie title in message
                processed_text = processed_text.replace("@" + movie_id, movie_title)

                response_knowledge.append(("movie_title", movie_title))
                #print("\n", movie_id)
                #print(conversation["movieMentions"].keys())
                mms = conversation["movieMentions"].get(movie_id)
                if isinstance(mms, dict):
                    mgenres = mms["genres"] if mms.get("genres") else [unk_terms["genre"]]
                    cast = [a["name"] + "," for a in mms["cast"]] if mms.get("cast") else [unk_terms["cast"]]
                    director = [a["name"] + "," for a in mms["director"]] if mms.get("director") else [unk_terms["director"]]

                    response_knowledge.extend([
                        ("movie_genre", " ".join(mgenres)),
                        ("director", " ".join(director)),
                        ("cast", " ".join(cast))
                    ])

            # For now, just pass in the surface form (later experiment is to try use normalized form)
            for genre_mention in message["genre_mentions"]:
                #print(genre_mention)
                response_knowledge.append(("genre", genre_mention["words"][0]))

            for imdb_entry in message["imdb_entries"]:
                soup = BeautifulSoup(imdb_entry, "xml")

                response_knowledge.append(('person', soup.find('name').text))


            if include_dact:
                dacts = []
                the_das = message["swbd_da"]
                for da in the_das:
                    d = "<" + da["label_name"].replace(" ", "_") + ">"
                    dact_set.add(d)
                    dacts.append(d)
                dact_tup = ("dact", " ".join(dacts))

                response_knowledge.append(dact_tup)
            #print(message)
            #print(response_knowledge)
            #input(">>")



            if i == len(messages) - 1 or \
                    message["senderWorkerId"] != messages[i + 1]["senderWorkerId"]:
                response += processed_text

                if not recommender_only or (recommender_only and sender_id != message["senderWorkerId"]):
                    examples.append(KnowledgeGroundedExample(
                        context,
                        response,
                        response_knowledge
                    ))

                    if len(response_knowledge) > 0:
                        num_examples_using_knowledge += 1

                context = context + [response]
                response = ""
                response_knowledge = []
            else:
                # We looked ahead and saw another follow-on response
                response += processed_text + " . "
    print("Num examples:", len(examples))
    print("Num examples using knowledge: ", num_examples_using_knowledge)
    return examples, list(unk_terms.values()) + list(dact_set)

def get_movie_db_map(mentions_file_path):
    movie_db_map = {}

    with open(mentions_file_path, 'r') as mentions_file:
        reader = csv.DictReader(mentions_file)

        for row in reader:
            movie_db_map[row['movieId']] = row['movieName']

    return movie_db_map

def try_load_pickle(pickle_file_path, get_special=False):
    print("trying to load pickle", pickle_file_path)
    if os.path.exists(pickle_file_path):
        with open(pickle_file_path, 'rb') as pickle_file:
            data = pickle.load(pickle_file)

        if get_special and isinstance(data, dict):
            retval = data["data"], data.get("special_terms")
        else:
            retval = data
        return retval
    print("not found ...")


def save_pickle(pickle_file_path, data, special_terms=None):

    if special_terms:
        data = {
            "data": data,
            "special_terms": special_terms
        }

    with open(pickle_file_path, 'wb') as pickle_file:
        pickle.dump(data, pickle_file)




def prepare_redial_baseline_dataset(
    redial_path,
    tokenizer,
    movie_db_map,
    dataset_cache_path='dataset_cache.pkl'
):

    dataset = try_load_pickle(dataset_cache_path)

    if dataset:
        print("Cached data already found, returning")
        return dataset

    split_files = {
        'train': 'train_data.jsonl',
        'test': 'test_data.jsonl'
    }

    dataset = {}

    for split, split_file_name in split_files.items():
        split_file_path = os.path.join(redial_path, split_file_name)
        examples = prepare_baseline_redial_split(split_file_path, tokenizer, movie_db_map)

        dataset[split] = examples

    save_pickle(dataset_cache_path, dataset)
    print("Saved file to cache ", dataset_cache_path)
    return dataset


def prepare_redial_knowledge_grounded_dataset(
    redial_path,
    tokenizer,
    movie_db_map,
    dataset_cache_path='kg_dataset_cache.pkl',
    split_files=None,
    recommender_only=False,
    include_dacts=True,
):
    dataset = try_load_pickle(dataset_cache_path, get_special=True)

    if dataset:
        print("Cached data already found, returning")
        return dataset[0], dataset[1]

    if split_files is None:
        split_files = {
            'train': 'train_data_genre_tagged.jsonl',
            'test': 'test_data_genre_tagged.jsonl'
        }

    dataset = {}

    for split, split_file_name in split_files.items():
        split_file_path = os.path.join(redial_path, split_file_name)
        examples, special_terms = prepare_redial_knowledge_grounded_split(split_file_path, movie_db_map, recommender_only, include_dacts)
        dataset[split] = examples
        if split.lower() == "train":
            train_terms = special_terms

    save_pickle(dataset_cache_path, dataset, special_terms=train_terms)
    print("Saved file to cache ", dataset_cache_path)
    return dataset, train_terms
<file_sep>certifi==2020.6.20
cffi # @ file:///tmp/build/80754af9/cffi_1600699165083/work
chardet==3.0.4
click==7.1.2
dataclasses==0.7
filelock==3.0.12
idna==2.10
IMDbPY==2020.9.25
joblib==0.17.0
lxml==4.6.1
mkl-fft==1.2.0
mkl-random==1.1.1
mkl-service==2.3.0
numpy # @ file:///tmp/build/80754af9/numpy_and_numpy_base_1603487797006/work
packaging==20.4
protobuf==3.13.0
pycparser # @ file:///tmp/build/80754af9/pycparser_1594388511720/work
pyparsing==2.4.7
regex==2020.10.28
requests==2.24.0
sacremoses==0.0.43
sentencepiece==0.1.91
six==1.15.0
SQLAlchemy==1.3.20
tokenizers==0.9.2
torch==1.3.1
tqdm==4.51.0
transformers==3.4.0
urllib3==1.25.11
<file_sep>"""
 Created by diesel
 12/19/19
"""


from __future__ import print_function, division


import os
import pandas as pd


def check_and_create(out_dir):
    if not os.path.exists(out_dir):
        print(" ** creating directory: {}".format(out_dir))
        os.makedirs(out_dir)


def write_lines(outpath, lines):
    with open(outpath, "w") as fout:
        fout.write("\n".join(lines))


def read_lines(fpath):
    with open(fpath, "r") as fin:
        lines = fin.read().split("\n")
    return lines


def read_csv_shards(dirpath, name_prefix, postfix):
    #print("reading csv shards")
    #print(" * dirpath:", dirpath)
    #print(" * name_prefix:", name_prefix)
    #print(" * postfix:", postfix)

    filenames = [fname for fname in os.listdir(dirpath) if fname.startswith(name_prefix) if fname.endswith(postfix)]
    #print(" * filenames:", filenames)

    idxs = [fname[len(name_prefix):-len(postfix)] for fname in filenames]
    #print(" * idxs:", idxs)
    sorted_names = [(int(idx), name) for idx, name in zip(idxs, filenames)]
    sorted_names.sort()

    data = pd.DataFrame()
    for idx, fname in sorted_names:
        df = pd.read_csv(os.path.join(dirpath, fname))
        data = pd.concat([data, df], axis=0)

    data.reset_index(inplace=True, drop=True)
    return data


def save_df_shards(to_write, outpath, max_size=15000):

    h = 0

    while to_write:
        h += 1
        new_df = pd.DataFrame(to_write[:max_size])
        outfile = outpath + "-{}.csv".format(h)
        if "Index" in new_df.columns.values:
            new_df.drop("Index", axis=1, inplace=True)
        new_df.to_csv(outfile, index=False, encoding="utf-8")
        to_write = to_write[max_size:]


def main():
    pass


if __name__ == "__main__":
    main()
<file_sep>
import pandas as pd
import numpy as np
import json
import nltk

def make_signature(row):
    return f"{row.original_title}  ({row.year})".lower()

def main1():
    """
    ['Headhunter  (2009)']
    :return:
    """
    print("hello there !")

    redial_df = pd.read_csv("movies_with_mentions.csv")
    redial_movies = {row.movieName.lower(): j for j, row in enumerate(redial_df.itertuples())}

    imdb_df = pd.read_csv("IMDb-movies.csv")
    imdb_df["sig"] = [make_signature(row) for row in imdb_df.itertuples()]
    #to_drop = [row.Index for row in imdb_df.itertuples() if row.sig not in redial_movies]
    #imdb_df.drop(to_drop, axis=0, inplace=True)
    imdb_movies = {row.sig: (row.imdb_title_id, row.title ) for j, row in enumerate(imdb_df.itertuples())}

    mapping = [imdb_movies.get(row.movieName.lower()) for row in redial_df.itertuples()]
    redial_df["imdb_index"] = mapping

    redial_df.to_csv("movies_with_mentions-IMDb.tsv", sep="\t", index=False)
    imdb_df.to_csv("IMDb-movies-redial.tsv", sep="\t", index=False)


def read_lines(fpath):
    with open(fpath, "r") as fin:
        return fin.read().split("\n")

_long_count = 0
_good_count = 0
def normalize_imbd_id(the_id, wanted_len):
    global _long_count, _good_count
    if not isinstance(the_id, str):
        the_id = str(the_id)
    if len(the_id) < wanted_len:
        # imdb movie id's are length 7 numeric identifiers with
        # leading zeroes when necessary.
        the_id = "0" * (wanted_len - len(the_id)) + the_id
    #assert len(the_id) == wanted_len, [the_id]
    if len(the_id) != wanted_len:
        _long_count += 1
    else:
        _good_count += 1
    return the_id


def get_movie(mid, movie_base):
    return movie_base.get(normalize_imbd_id(mid, 7))

def get_person(pid, person_base):
    if isinstance(pid, str) or isinstance(pid, int):
        retval = person_base.get(normalize_imbd_id(pid, 7))
    else:
        retval = pid
    return retval

def main():

    movielense_2_imdb = {}
    database_2_imdb = {}
    df = pd.read_csv("movies_merged_with_imdb.csv")
    for row in df.itertuples():
        #print(row)
        imdbId = normalize_imbd_id(row.imdbId, 7)
        movielense_2_imdb[row.movielensId] = imdbId
        database_2_imdb[str(row.databaseId)] = imdbId
        assert isinstance(imdbId, str), imdbId

    _imdb_movie_set = set(movielense_2_imdb.values())
    assert "3874544" in _imdb_movie_set
    assert "0247745" in _imdb_movie_set
    assert database_2_imdb["111776"] == "0247745"

    movies_db, people_db = load_databases()
    #exit(0)
    for split_name in ["test", "train"]:
        new_lines = []
        fname = f"{split_name}_data_swda_tagged.jsonl"
        print("loading", fname)
        print("processing ...")
        for line in read_lines(fname):
            try:
                d = json.loads(line)
            except Exception as e:
                # copy the line over
                new_lines.append(line)
                # and continue to the next line
                continue

            new_mentions = {}
            for mid in d["movieMentions"]:

                #print("\nmid:", mid)
                imdb_id = database_2_imdb.get(mid)
                #print("imdb_id:", imdb_id)
                #print("mention:", d["movieMentions"][mid])
                movie = get_movie(imdb_id, movies_db)
                #print("movie:", movie)


                if isinstance(movie, dict):
                    for attr_name in ["director", "cast"]:
                        new_vals = [get_person(pid, people_db) if get_person(pid, people_db) else pid
                                    for pid in movie.get(attr_name, [])]

                        for p in new_vals:
                            if not isinstance(p, dict):
                                continue
                            new_kf_movies = []
                            for kf_movie_id in p.get("known for", []):
                                #print("kf_movie_id:", [kf_movie_id])
                                if isinstance(kf_movie_id, dict):
                                    kf_movie = kf_movie_id
                                else:
                                    kf_movie = get_movie(kf_movie_id, movies_db)

                                #print("kf_movie:", kf_movie)
                                if kf_movie is None:
                                    new_kf_movies.append(kf_movie_id)
                                else:
                                    new_kf_movies.append({
                                        "long imdb title": kf_movie.get("long imdb title", "_unk_title_"),
                                        "imdb_id": kf_movie_id.get("imdb_id") if isinstance(kf_movie_id, dict) else kf_movie_id
                                    })
                            p["known for"] = new_kf_movies
                        movie[attr_name] = new_vals
                    movie["imdb_id"] = imdb_id
                    new_mentions[mid] = movie

                else:
                    # blindly copy whatever is there.
                    new_mentions[mid] = d["movieMentions"][mid]
            #print(new_mentions[mid])
            d["movieMentions"] = new_mentions



            new_lines.append(json.dumps(d))

            if input(">>>"):
                exit(0)


        fname = f"{split_name}-with-movie-info.jsonl"
        print("saving", fname)
        with open(fname, "w") as fout:
            fout.write("\n".join(new_lines))

        print("_long_count:", _long_count)
        print("_good_count:", _good_count)


def load_databases():
    with open("all-movies.json", "r") as fin:
        movies_db = json.load(fin)
    mkeys = [len(str(mid)) for mid in list(movies_db.keys())]
    print("max movie key length:", max(mkeys))
    print("min movie key length:", min(mkeys))
    movies_db = {normalize_imbd_id(imdb_id, 7): entry
                   for imdb_id, entry in movies_db.items()}


    with open("all-people.json", "r") as fin:
        people_db = json.load(fin)

    pkeys = [len(str(mid)) for mid in list(movies_db.keys())]
    print("max people key length:", max(pkeys))
    print("min people key length:", min(pkeys))
    people_db = {normalize_imbd_id(imdb_id, 7): entry
                   for imdb_id, entry in people_db.items()}

    return movies_db, people_db


if __name__ == "__main__":
    main()



