Dependency Parsing

      Wyobraź                          
   ______|_____________                 
  |      |         człowieka           
  |      |             |                
  |      |        znajdującego         
  |      |      _______|__________      
  |      |     |                  na   
  |      |     |                  |     
  |      |     |                górze  
  |      |     |                  |     
sobie    .    się              ogromnej

python3 validation_script.py --train
python3 validation_script.py

def points(root_placement, uuas):
    def scale(x, lower=0.5, upper=0.85):
        scaled = min(max(x, lower), upper)
        return (scaled - lower) / (upper - lower)
    return (scale(root_placement) + scale(uuas))

FINAL_EVALUATION_MODE = False  # While checking your solution, we will change this value to True
DEPTH_MODEL_PATH = 'depth_model.pth'  # Do not change!
DISTANCE_MODEL_PATH = 'distance_model.pth'  # Do not change!

from typing import List

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (AutoModel, AutoTokenizer, PreTrainedModel,
                          PreTrainedTokenizer)
from utils import (ListDataset, ParsedSentence, Sentence, merge_subword_tokens,
                   read_conll, uuas_score)

tokenizer = AutoTokenizer.from_pretrained("allegro/herbert-base-cased")
model = AutoModel.from_pretrained("allegro/herbert-base-cased")

train_sentences = read_conll('train.conll')  # 1000 sentences
val_sentences = read_conll('valid.conll')  # 200 sentences

train_sentences[6].pretty_print()  # display the tree of one sentence
print(train_sentences[6])

def get_distances(sentence: ParsedSentence):
    """Find the distances between every pair of words in the sentence.
    Returns a numpy matrix of shape (len(sentence), len(sentence))."""

    # TODO: implement me
    ...

    return distances

print(get_distances(train_sentences[1]))

def get_bert_embeddings(
    sentences_s: List[str],
    tokenizer: PreTrainedTokenizer, 
    model: PreTrainedModel,
    progress_bar: bool = False,
):
    """
    The function returns subword embeddings for a list of sentences.

    Arguments:
        sentences_s: List of sentences. Each sentence is represented as a string.
        tokenizer: HERBERT tokenizer
        model: HERBERT model
        progress_bar: Whether to display a progress bar.

    Returns:
        tokens: A list that, for each sentence, contains the list of subword tokens of that sentence.
        embeddings: A list that, for each sentence, contains a list of tensors of shape 
            (seq_len, emb_dim). Note that seq_len may differ between sentences.
    """

    # Hints:
    #  1. You can use the functions:
    #   encoded = tokenizer.batch_encode_plus(...)
    #   with torch.no_grad():
    #     model(**encoded, output_hidden_states=True)
    #  2. To speed up the computation, remember to group (batch) the sentences before passing them to the model.
    #  3. Remember that each sentence may have a different length, so HERBERT will apply padding to fill the extra
    #   space in the returned tensor. Remember to remove the padding from the results.
    #  4. The tokenizer and the model use special tokens (e.g. for the beginning and end of a sentence), which should also 
    #   be removed.

    # TODO: implement me
    ...

    return tokens, embeddings

def get_word_embeddings(sentences: List[Sentence], tokenizer, model):
    """The function returns word embeddings for a list of sentences, using the model and the tokenizer."""

    # Hints:
    #  1. Use the function get_bert_embeddings to obtain the subword embeddings.
    #  2. Use the function merge_subword_tokens to obtain the word embeddings.

    # TODO: implement me
    ...

    return embeddings

def get_datasets(sentences: List[ParsedSentence], tokenizer, model):
    embeddings = get_word_embeddings(sentences, tokenizer, model)
    distances = [get_distances(sent) for sent in sentences]
    depths = [dist[sent.root][..., None] for dist, sent in zip(distances, sentences)]
    dataset_dist = ListDataset(list(zip(embeddings, distances, sentences)))
    dataset_depth = ListDataset(list(zip(embeddings, depths, sentences)))
    return dataset_dist, dataset_depth


if not FINAL_EVALUATION_MODE:
    trainset_dist, trainset_depth =  get_datasets(train_sentences, tokenizer, model)
    valset_dist, valset_depth = get_datasets(val_sentences, tokenizer, model)

def pad_arrays(sequence, pad_with=np.inf):
    """
    Assumes that sequence contains arrays (ndarrays) with the same number of dimensions.
    Returns a tensor containing the data padded to the same dimensions with the value pad_with, 
    where the sequence index corresponds to the first dimension.
    """

    shapes = np.array([list(seq.shape) for seq in sequence])
    max_lens = list(shapes.max(axis=0))
    padded = [np.pad(
                seq, 
                tuple((0, max_lens[i] - seq.shape[i]) for i in range(seq.ndim)), 
                'constant', 
                constant_values=pad_with
            ) for seq in sequence]
    return torch.tensor(padded)


def collate_fn(batch):
    embeddings, targets, sentences = zip(*batch)
    padded_embeddings = pad_arrays(embeddings, pad_with=0)
    padded_targets = pad_arrays(targets, pad_with=np.inf)
    mask = padded_targets != torch.inf
    return padded_embeddings, padded_targets, mask, sentences


if not FINAL_EVALUATION_MODE:
    dist_trainloader = DataLoader(trainset_dist, batch_size=32, shuffle=True, collate_fn=collate_fn)
    dist_valloader = DataLoader(valset_dist, batch_size=32, shuffle=False, collate_fn=collate_fn)

    depth_trainloader = DataLoader(trainset_depth, batch_size=32, shuffle=True, collate_fn=collate_fn)
    depth_valloader = DataLoader(valset_depth, batch_size=32, shuffle=False, collate_fn=collate_fn)

# dist_trainloader and dist_valloader return tuples (embeddings, distances, masks, sentences)
# depths_trainloader and depths_valloader return tuples (embeddings, depths, masks, sentences)  
# embeddings.shape: (batch_size, max_seq_len, emb_dim)
# distances.shape: (batch_size, max_seq_len, max_seq_len)
# depths.shape: (batch_size, max_seq_len, 1)

class DistanceModel(torch.nn.Module):
    def __init__(self):
        # TODO: implement me
        ...

    def forward(self, x):
        # TODO: implement me
        ...


class DepthModel(torch.nn.Module):
    def __init__(self):
        # TODO: implement me
        ...

    def forward(self, x):
        # TODO: implement me
        ...

def loss_fn(output, target, mask):
    # mask, target and output are tensors of the same shape
    # mask contains 1 where target contains data, and 0 where there is padding

    # TODO: implement me
    ...


def train_model(model, dataloader, valloader, epochs, lr):
    """Training loop for your models."""
    # TODO: implement me
    ...


# During evaluation, the models should not be trained again.
if not FINAL_EVALUATION_MODE: 
    print("Training depth model")
    depth_model = DepthModel()
    # TODO: set the hyperparameters
    train_model(depth_model, depth_trainloader, depth_valloader, lr=..., epochs=...)  
    # save the model weights to a file
    torch.save(depth_model.state_dict(), DEPTH_MODEL_PATH)

    print("Training distance model")
    distance_model = DistanceModel()
    # TODO: set the hyperparameters
    train_model(distance_model, dist_trainloader, dist_valloader, lr=..., epochs=...)
    # save the model weights to a file
    torch.save(distance_model.state_dict(), DISTANCE_MODEL_PATH)

def parse_sentence(sent: Sentence, distance_model, depth_model, tokenizer, model) -> ParsedSentence:
    """Build the syntax tree for a single sentence.

    Arguments:
        sent: The sentence to parse.
        distance_model: Trained distance model
        depth_model: Trained depth model
        tokenizer: HERBERT tokenizer
        model: HERBERT model

    Returns:
        ParsedSentence: The sentence with the predicted syntax tree.

    """

    # Your solution should:
    # 1. Obtain the word embeddings for the sentence.
    # 2. Choose the root of the syntax tree heuristically, using depth_model.
    # 3. Compute the distances between every pair of nodes, using distance_model.
    # 4. Implement a heuristic method of your own design for choosing the edges 
    #    of the tree on the basis of the predicted distances.
    # 5. Obtain a ParsedSentence object. You can use the function ParsedSentence.from_edges_and_root
    # 6. Return the ParsedSentence object.

    # Hints:
    #  You can use sent.pretty_print() to visualise the parsed sentence.

    # Note:
    # This function will be used to score your solution. This function should return an actual tree,
    # with len(sent) - 1 edges. If your prediction is not a tree, it will be invalid and
    # you will not get any points for it. If you want to get only some of the points, by competing only in the 
    # root placement metric, you should still return a valid tree.

    # TODO: implement me
    ...

if not FINAL_EVALUATION_MODE:
    sent = train_sentences[30]
    parse_sentence(sent, distance_model, depth_model, tokenizer, model).pretty_print()  # Predicted tree
    sent.pretty_print()  # Gold tree (from the dataset)
    print(sent)

def points(root_placement, uuas):
    def scale(x, lower=0.5, upper=0.85):
        scaled = min(max(x, lower), upper)
        return (scaled - lower) / (upper - lower)
    return (scale(root_placement) + scale(uuas))

def evaluate_model(sentences: List[ParsedSentence], distance_model, depth_model, tokenizer, model):
    sum_uuas = 0
    root_correct = 0
    with torch.no_grad():
        for sent in sentences:
            parsed = parse_sentence(sent, distance_model, depth_model, tokenizer, model)
            root_correct += int(parsed.root == sent.root)
            sum_uuas += uuas_score(sent, parsed)
    
    root_placement = root_correct / len(sentences)
    uuas = sum_uuas / len(sentences)

    print(f"UUAS: {uuas * 100:.3}%")
    print(f"Root placement: {root_placement * 100:.3}%")
    print(f"Your score: {points(root_placement, uuas):.1}/2.0")

if not FINAL_EVALUATION_MODE:
    distance_model_loaded = DistanceModel()
    distance_model_loaded.load_state_dict(torch.load(DISTANCE_MODEL_PATH))

    depth_model_loaded = DepthModel()
    depth_model_loaded.load_state_dict(torch.load(DEPTH_MODEL_PATH))

    evaluate_model(val_sentences, distance_model_loaded, depth_model_loaded, tokenizer, model)