text: '''
  Usage:
    quote_detection.py <model-name> -t <corpus-type> -c <corpus-path> -e <embedding-path>
  
  Options:
    -t           Corpus type (either parc, rwg, or stop)
    -c           Path to corpus
    -e           Path to embeddings file
  
  '''
  
  
  
  import random
  from docopt import docopt
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  from torch.nn.parameter import Parameter
  import numpy as np
  
  from evaluate import evaluate, report
  from progressify import progressify
  
  class LSTMSeq2Seq(nn.Module):
      def __init__(self, embedding_dim, hidden_dim,  n_labels, layers=2, bidirectional=True):
          super().__init__()
          self.embedding_dim = embedding_dim
          self.hidden_dim = hidden_dim
          self.n_labels = n_labels
          self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=layers, bidirectional=bidirectional, dropout=0.5, batch_first=True)
          self.linear = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, n_labels)
          
      def forward(self, x):
          x = self.lstm(x)[0]
          if isinstance(x, nn.utils.rnn.PackedSequence):
              data = self.linear(x.data)
              return nn.utils.rnn.PackedSequence(data, x.batch_sizes)
          else:
              x = self.linear(x)
              return x
  
  class SelfAttentionLayer(nn.Module):
      def __init__(self, d_x, d_k, d_v):
          super().__init__()
          self.d_x = d_x
          self.d_k = d_k
          self.d_v = d_v
          self.w_q = nn.Linear(d_x, d_k)
          self.w_k = nn.Linear(d_x, d_k)
          self.w_v = nn.Linear(d_x, d_v)
  
      def forward(self, x):
          # x: float[batch, sequence_length, d_x]
          Q = self.w_q(x)
          # Q: float[batch, sequence_length, d_k]
          K = self.w_k(x)
          # K: float[batch, sequence_length, d_k]
          V = self.w_v(x)
          # V: float[batch, sequence_length, d_v]
          logits = torch.bmm(K, V.permute(0, 2, 1)) / np.sqrt(self.d_k)
          # logits float[batch, sequence_length, sequence_length]
          return torch.bmm(torch.softmax(logits, dim=-1), V)
          # return float[batch, sequence_length, d_v]
  
  class MultiHeadedAttentionLayer(nn.Module):
      def __init__(self, d_x, n_heads):
          super().__init__()
          assert d_x % n_heads == 0
          self.n_heads = n_heads
          self.heads = [
              SelfAttentionLayer(d_x, d_x // n_heads, d_x // n_heads)
              for _ in range(n_heads)
          ]
  
      def forward(self, x):
          # x: float[batch, sequence_length, d_x]
          return torch.cat([
              head(x) for head in self.heads
          ], dim=-1)
  
  
  class TransformerLayer(nn.Module):
      def __init__(self, d_x, n_heads, activation=F.relu):
          super().__init__()
          self.d_x = d_x
          self.n_heads = n_heads
          self.activation = activation
          self.attention = MultiHeadedAttentionLayer(d_x, n_heads)
          self.linear = nn.Linear(d_x, d_x)
          self.ln1 = nn.LayerNorm([d_x])
          self.ln2 = nn.LayerNorm([d_x])
          self.dropout = nn.Dropout(p=0.1)
  
      def forward(self, x):
          # x: float[batch, sequence, d_x]
          x_attn = self.dropout(self.attention(x))
          # x_attn: float[batch, sequence, d_x]
          x = self.ln1(x + x_attn)
          x_ff = self.dropout(self.activation(self.linear(x)))
          x = self.ln2(x + x_ff)
          return x
          
          
  class TransformerSeq2Seq(nn.Module):
      def __init__(self, d_x, n_layers, n_heads, d_out, pos_encode=False):
          super().__init__()
          self.d_x = d_x
          self.d_out = d_out
          self.n_heads = n_heads
          self.pos_encode = pos_encode
          self.layers = [TransformerLayer(d_x, n_heads) for _ in range(n_layers)]
          self.out = nn.Linear(d_x, d_out)
          self.dropout = nn.Dropout(p=0.1)
  
      def _pos_encode(self, x):
          # x: float[batch, sequence_length, n_dims]
          batch, sequence_length, n_dims = x.shape
          positions = torch.arange(sequence_length, dtype=torch.float)
          # positions: float[sequence_length]
          frequencies = 10000 ** (-torch.arange(n_dims/2, dtype=torch.float) / (n_dims/2))
          # frequencies: float[n_dims / 2]
          coss = torch.cos(torch.ger(positions, frequencies))
          sins = torch.sin(torch.ger(positions, frequencies))
          pes = torch.cat([coss, sins], -1)
          # pes: float[sequence_length, n_dims]
          return self.dropout(x + pes)
          
      def forward(self, x):
          was_packed = False
          if isinstance(x, nn.utils.rnn.PackedSequence):
              was_packed = True
              x, lengths = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
          # x float:[batch, sequence_length, d_x]
          if self.pos_encode:
              x = self._pos_encode(x)
          for layer in self.layers:
              x = layer(x)
          x = self.out(x)
          if was_packed:
              x = nn.utils.rnn.pack_padded_sequence(x, lengths, batch_first=True)
          return x
  
  class ComposedSeq2Seq(nn.Module):
      def __init__(self, dim, lstm_layers=2, transformer_layers=6):
          super().__init__()
          self.lstm = LSTMSeq2Seq(dim, dim, dim, layers=lstm_layers)
          self.transformer = TransformerSeq2Seq(dim, transformer_layers, 8, dim)
  
      def forward(self, x):
          return self.lstm(self.transformer(x))
          
  class IdentityLayer(nn.Module):
      def forward(self, x):
          return x
  
  class LSTMCRF(nn.Module):
      # only supports packedsequences of length 1...
      def __init__(self, dim, n_tags, lstm_layers=2):
          super().__init__()
          self.lstm = LSTMSeq2Seq(dim, dim, dim, layers=lstm_layers)
          self.crf = CRFLayer(dim, n_tags)
          # hack so we can print these nicely
          self.transitions = self.crf.transitions
  
      def forward(self, x):
          x = self.lstm(x)
          if isinstance(x, nn.utils.rnn.PackedSequence):
              return self.crf(x.data)
          else:
              return self.crf(x)
  
      def neg_log_likelihood(self, x, tags):
          x = self.lstm(x)
          if isinstance(x, nn.utils.rnn.PackedSequence):
              return self.crf.neg_log_likelihood(x.data, tags.data)
          else:
              return self.crf.neg_log_likelihood(x, tags)
  
  
  class QuoteDetectionModel(nn.Module):
      def __init__(
          self,
          n_features,
          dim=300,
          lam=1e-4,
          encoder=None,
          use_crf=False,
          transformer=False,
          sample_steps=1,
          label_scheme='BE',
          viterbi=False,
          pipeline=True,
      ):
          super().__init__()
          self.n_features = n_features
          self.dim = dim
          self.lam = lam
          self.use_crf = use_crf
          self.sample_steps = sample_steps
          self.label_scheme = label_scheme
          self.viterbi = viterbi
          self.pipeline = pipeline
          self.embedding = nn.EmbeddingBag(n_features, dim, mode='sum')
          self.roles = list(roles)
          if encoder is None:
              self.encoder = lambda x: x
          elif encoder == 'transformer':
              self.encoder = TransformerSeq2Seq(dim, 6, 10, dim, pos_encode=True)
          if use_crf:
              self.seq2seqs = nn.ModuleList()
              self.crf_encoders = nn.ModuleList()
              self.crfs = nn.ModuleList()
              for role in self.roles:
                  if transformer:
                      self.seq2seqs.append(TransformerSeq2Seq(dim, 6, 10, dim, pos_encode=True))
                  else:
                      self.seq2seqs.append(LSTMSeq2Seq(dim, dim, dim, layers=2, bidirectional=True))
                  #self.crfs.append(CRFLayer(dim, 3))
                  self.crf_encoders.append(nn.Linear(dim, 3))
                  self.crfs.append(ConditionalRandomField(3))
          else:
              self.seq2seqs = nn.ModuleList()
              self.outputs = nn.ModuleList()
              for role in self.roles:
                  if transformer:
                      self.seq2seqs.append(TransformerSeq2Seq(dim, 6, 10, dim, pos_encode=True))
                  else:
                      self.seq2seqs.append(LSTMSeq2Seq(dim, dim, dim, layers=2, bidirectional=True))
                  self.outputs.append(nn.Linear(dim, 3))
          self.tag_embeddings = Parameter(torch.Tensor(3, 3, dim))
          nn.init.normal_(self.tag_embeddings)
  
      def forward(self, x):
          # x: long[batch_size, sequence_length, bag_size]
          nn.init.zeros_(self.embedding.weight[0])
          batch_size, sequence_length, bag_size = x.shape
          embedded = self.embedding(x.view(-1, bag_size)).view(batch_size, sequence_length, self.dim)
          # embedded: float[batch_size, sequence_length, self.dim]
          embedded = self.encoder(embedded)
          prediction_embeddings = torch.zeros_like(embedded)
          if self.use_crf:
              for step in range(self.sample_steps):
                  paths = []
                  for i, (seq2seq, crf_encoder, crf) in enumerate(zip(self.seq2seqs, self.crf_encoders, self.crfs)):
                      res = seq2seq(embedded + prediction_embeddings)
                      # res: float[batch_size, sequence_length, self.dim]
                      sequence_length = res.shape[1]
                      path, score = crf.viterbi_tags(crf_encoder(res).cpu(), torch.tensor([sequence_length], dtype=torch.int64).cpu())[0]
                      path = torch.tensor([path])
                      if self.pipeline:
                          prediction_embeddings = prediction_embeddings + F.embedding(path, self.tag_embeddings[i])
                      paths.append(path)
              return torch.stack(paths, dim=1)
              # return float:[batch_size, len(self.roles), sequence_length]
          else:
              for step in range(self.sample_steps):
                  paths = []
                  for i, (seq2seq, output) in enumerate(zip(self.seq2seqs, self.outputs)):
                      res = seq2seq(embedded + prediction_embeddings)
                      # res: float[batch_size, sequence_length, self.dim]
                      logits = output(res)
                      # logits: float[batch_size, sequence_length, 3]
                      if self.viterbi:
                          #breakpoint()
                          probs = torch.softmax(logits, dim=-1)
                          path = self.viterbi_sequence(probs)
                      else:
                          path = torch.argmax(logits, dim=-1)
                      if self.pipeline:
                          prediction_embeddings = prediction_embeddings + F.embedding(path, self.tag_embeddings[i])
                      paths.append(path)
              return torch.stack(paths, dim=1)
              # return float:[batch_size, len(self.roles), sequence_length]
              
      def neg_log_likelihood(self, x, tags):
          # x: long[batch_size, sequence_length, bag_size]
          # tags: int[batch_size, len(self.roles), sequence_length]
          nn.init.zeros_(self.embedding.weight[0])
          batch_size, sequence_length, bag_size = x.shape
          embedded = self.embedding(x.view(-1, bag_size)).view(batch_size, sequence_length, self.dim)
          # embedded: float[batch_size, sequence_length, self.dim]
          embedded = self.encoder(embedded)
          nll = 0
          prediction_embeddings = torch.zeros_like(embedded)
          if self.use_crf:
              for step in range(self.sample_steps):
                  for i, (seq2seq, crf_encoder, crf) in enumerate(zip(self.seq2seqs, self.crf_encoders, self.crfs)):
                      role_tags = tags[:,i,:]
                      res = seq2seq(embedded + prediction_embeddings)
                      nll -= torch.sum(crf(crf_encoder(res), role_tags))
                      if self.pipeline:
                          score, path = crf(res)
                          prediction_embeddings = prediction_embeddings + F.embedding(path, self.tag_embeddings[i])
              return nll
          else:
              for step in range(self.sample_steps):
                  for i, (seq2seq, output) in enumerate(zip(self.seq2seqs, self.outputs)):
                      role_tags = tags[:,i,:]
                      # role_tags: long[batch_size, sequence_length]
                      res = seq2seq(embedded + prediction_embeddings)
                      # res: float[batch_size, sequence_length, self.dim]
                      logits = output(res)
                      # logits: float[batch_size, sequence_length, 3]
                      nll += F.cross_entropy(logits.view(-1, 3), role_tags.view(-1), reduction='sum')
                      path = torch.argmax(logits, dim=-1)
                      if self.pipeline:
                          if self.viterbi:
                              # It is too expensive to do the viterbi step during training I think
                              # instead, let's use gold-standard labels
                              # error propogation yada yada
                              prediction_embeddings = prediction_embeddings + F.embedding(role_tags, self.tag_embeddings[i])
                          else:
                              prediction_embeddings = prediction_embeddings + F.embedding(path, self.tag_embeddings[i])
              return nll
  
      def viterbi_sequence(self, probs):
          if self.label_scheme == "BE":
              #naive = [" BE"[m] for m in torch.argmax(probs, -1)]
              paths = torch.zeros(probs.shape[:-1], dtype=torch.long)
              inf = float('inf')
              # OIB probabilities
              for i, pi in enumerate(probs):
                  state_lps = [(0, -inf, -inf)]
                  for lp_, lpb, lpe in torch.log(pi):
                      state_lpb = max(state_lps[-1]) + lpb
                      state_lpi = max(state_lps[-1][1], state_lps[-1][2]) + lp_
                      state_lpo = max(
                          state_lps[-1][0] + lp_,
                          state_lps[-1][1] + lpe,
                          state_lps[-1][2] + lpe
                      )
                      state_lps.append((state_lpo, state_lpi, state_lpb))
  
  
                  #construct the most likely BIO sequence from back to front
                  bios = []
                  outside = True
                  for (lpo, lpi, lpb) in reversed(state_lps[1:]):
                      if outside:
                          if lpo > lpi and lpo > lpb:
                              bios.append('O')
                          elif lpb > lpo and lpb > lpi:
                              bios.append('B')
                          else:
                              bios.append('I')
                              outside = False
                      else:
                          if lpb > lpi:
                              bios.append('B')
                              outside = True
                          else:
                              bios.append('I')
  
                  bios.reverse()
  
                  outside = True
                  for j, bio in enumerate(bios):
                      if bio == 'B':
                          #final_labels.append('B')
                          paths[i,j] = 1
                          outside = False
                      elif bio == 'I':
                          #final_labels.append(' ')
                          paths[i,j] = 0
                          assert not outside
                      elif bio == 'O':
                          if not outside:
                              #final_labels.append('E')
                              paths[i,j] = 2
                          else:
                              #final_labels.append(' ')
                              paths[i,j] = 0
                          outside = True
          return paths
          
  def jointshuffle(l1, l2):
      print("shuffling...")
      zipped = list(zip(l1, l2))
      random.shuffle(zipped)
      c, d = zip(*zipped)
      print("done shuffling")
      return list(c), list(d)
  
  
  def batchify(feats, labels, batch_size):
      feats, labels = jointshuffle(feats, labels)
      feats.sort(key=len)
      labels.sort(key=lambda lab:lab.shape[1])
      
  
  def encode_feats(*corpora):
      seen_features = {None: 0}
  
      encoded_corpora = []
      for corpus in corpora:
          encoded_corpus = []
          for document in corpus:
              bag_size = max([len(token) for token in document])
              document_tensor = []
              for token in document:
                  bag = []
                  for feature in token:
                      if feature not in seen_features:
                          seen_features[feature] = len(seen_features)
                      bag.append(seen_features[feature])
                  while len(bag) < bag_size:
                      bag.append(0)
                  document_tensor.append(bag)
              document_tensor = torch.tensor(document_tensor, dtype=torch.long)
              encoded_corpus.append(document_tensor)
          encoded_corpora.append(encoded_corpus)
      return encoded_corpora, seen_features
  
  def encode_labels(*corpora, scheme='BE'):
      if scheme == 'BE':
          tags = ' BE'
      elif scheme == 'BIO':
          tags = 'OIB'
  
      encoded_corpora = []
      for corpus in corpora:
          encoded_corpus = []
          for document in corpus:
              document_tensor = []
              for role in roles:
                  role_tensor = []
                  for token in document:
                      role_tensor.append(tags.index(token[role]))
                  document_tensor.append(role_tensor)
              document_tensor = torch.tensor(document_tensor, dtype=torch.long)
              encoded_corpus.append(document_tensor)
          encoded_corpora.append(encoded_corpus)
      return encoded_corpora
  
  def inject_pretrained_embeddings(model, embedding_path, feat_indices):
      print("loading pre-trained embeddings...")
      with open(embedding_path) as f:
          for line in f:
              line = line.split(' ')
              word = line[0]
              if ('word', word) in feat_indices:
                  index = feat_indices[('word', word)]
                  v = torch.Tensor([float(li) for li in line[1:]])
                  model.embedding.weight.data[index] = v
              if ('lemma', word) in feat_indices:
                  index = feat_indices[('lemma', word)]
                  v = torch.Tensor([float(li) for li in reversed(line[1:])])
                  model.embedding.weight.data[index] = v
  
      print("done!")
  def eval(loss_func, feats, labels):
      with torch.no_grad():
          loss = 0
          for fi, li in zip(feats, progressify(labels, "Evaluating datum %%i / %d" % len(labels))):
              loss += loss_func(fi, li)
          return loss / len(feats)
  
  def train(loss_func, optimizer, feats, labels, lamb=1e-4):
      feats, labels = jointshuffle(feats, labels)
      mean_loss = None
      def progressify_str(i, _):
          s = "training datum %d / %d." % (i, len(labels))
          if i > 0:
              s += " Mean training loss: %f" % mean_loss
          return s
      for fi, li in zip(feats, progressify(labels, progressify_str)):
          optimizer.zero_grad()
          f = fi.unsqueeze(0)
          loss = loss_func(f, li.unsqueeze(0))
          if mean_loss is None:
              mean_loss = loss
          else:
              mean_loss = .995 * mean_loss + .005 * loss
          
          # l2 regularization
          for param_group in optimizer.param_groups:
              for param in param_group['params']:
                  loss += lamb * torch.sum(param ** 2)
          loss.backward()
          optimizer.step()
  
  def predict(forward_func, feats):
      predictions = []
      with torch.no_grad():
          for datum in feats:
              batch = datum.unsqueeze(0)
              #batch_feats = torch.stack(feats[i:i+predict_batch_size])
              batch_predictions = forward_func(batch)
              for pi in batch_predictions:
                  predictions.append(pi)
      return predictions
  
  def train_loop(model, optimizer, train_feats, train_labels, dev_feats, dev_labels, gamma=0.75, callback=None):
      loss_func = model.neg_log_likelihood
      #breakpoint()
      running_average = -1
      epoch = 0
      while True:
          print("Epoch %d" % epoch)
          model.eval()
          dev_score = callback()
          #print("Dev loss: %f" % dev_loss)
          running_average = gamma * running_average + (1-gamma) * dev_score
          print("Running average: %f" % running_average)
          if dev_score < running_average:
              break
          model.train()
          train(loss_func, optimizer, train_feats, train_labels)
          epoch += 1
      
  
  def get_ev(model, feats, raw_labels, eval_mode='exact'):
      print("Predicting spans...")
      predicted = predict(model, feats)
      predicted_processed = []
      for doc in predicted:
          doc_processed = []
          for i in range(len(doc[0])):
              token_processed = {}
              for r, role in enumerate(roles):
                  if scheme == 'BE':
                      token_processed[role] = ' BE  '[doc[r][i]]
                  elif scheme == 'BIO':
                      token_processed[role] = 'OIBOO'[doc[r][i]]
              doc_processed.append(token_processed)
          predicted_processed.append(doc_processed)
      # BUG HERE!! should say scheme=scheme
      return evaluate(predicted_processed, raw_labels, roles=roles, mode=eval_mode)
  
  
  
  
  def run_model(
          raw_train_feats, raw_train_labels, raw_dev_feats, raw_dev_labels, raw_test_feats, raw_test_label
  ):
      (train_feats, dev_feats, test_feats), feat_indices = encode_feats(raw_train_feats, raw_dev_feats, raw_test_feats)
      train_labels, dev_labels, test_labels = encode_labels(raw_train_labels, raw_dev_labels, raw_test_labels)
      
      n_feats = len(feat_indices)
  
      model = QuoteDetectionModel(n_feats, use_crf=use_crf, sample_steps=1, label_scheme=scheme, viterbi=False, transformer=False, pipeline=False)
  
      if embedding_path is not None:
          inject_pretrained_embeddings(model, embedding_path, feat_indices)
  
      optimizer = optim.Adam(model.parameters())
      best_f1 = -1
      def training_callback():
          nonlocal best_f1
          if check_presence:
              ev = get_ev(model, dev_feats, raw_dev_labels, eval_mode='presence')
          else:
              ev = get_ev(model, dev_feats, raw_dev_labels)
          print(report(ev, roles=roles))
          f1 = 0
          if 'content' in ev:
              tp = ev['content']['tp']
              fp = ev['content']['fp']
              fn = ev['content']['fn']
          else:
              tp = 0
              fp = 0
              fn = 0
              for role in ev:
                  tp += ev[role]['tp']
                  fp += ev[role]['fp']
                  fn += ev[role]['fn']
                  
          if tp != 0:
              p = tp / (tp + fp)
              r = tp / (tp + fn)
              f1 = 2 / (1/p + 1/r)
          if f1 > best_f1:
              print('Best model so far! Saving...')
              best_f1 = f1
              torch.save(model.state_dict(), model_path)
          return f1
  
      train_loop(model, optimizer, train_feats, train_labels, dev_feats, dev_labels, callback=training_callback)
      print("Loading best model...")
      model.load_state_dict(torch.load(model_path))
      print("Evaluating on test-set...")
      ev = get_ev(model, test_feats, raw_test_labels, eval_mode='exact')
      print(report(ev, roles=roles))
      if check_presence:
          ev_presence = get_ev(model, test_feats, raw_test_labels, eval_mode='presence')
          print('presence/absence:')
          print(report(ev_presence, roles=roles))
          return ev, ev_presence
      else:
          return ev
      
  
  if __name__ == '__main__':
      arguments = docopt(__doc__)
      model_name = arguments['<model-name>']
      corpus_type = arguments['<corpus-type>']
      corpus_path = arguments['<corpus-path>']
      embedding_path = arguments['<embedding-path>']
  
      assert corpus_type in {'parc', 'stop', 'rwg'}
  
      xvalidate = (corpus_type in {'stop', 'rwg'})
      check_presence = (corpus_type == 'rwg')
  
  
      if corpus_type in {'rwg', 'stop'}:
          roles = [
              'direct',
              'indirect',
              'free_indirect',
              'reported'
          ]
      elif corpus_type == 'parc':
          roles = ['content']
  
      cuda = True
      if cuda:
          torch.set_default_tensor_type('torch.cuda.FloatTensor')
      
      if corpus_type == 'rwg':
          from rwg2feat import corpus_feats_and_labels
      elif corpus_type == 'stop':
          from stop2feat import corpus_feats_and_labels
      elif corpus_type == 'parc':
          from parc2feat import corpus_feats_and_labels
  
      import sys
      import os
  
      use_crf = False
      scheme = 'BE'
  
      model_path = model_name + '.pkl'
  
      if xvalidate:
          raw_feats, raw_labels = corpus_feats_and_labels(corpus_path, label_scheme=scheme)
          i_dev = len(raw_feats) // 10
          i_train = 2 * i_dev
          raw_feats, raw_labels = jointshuffle(raw_feats, raw_labels)
          for step in range(10):
              print("Cross-validation step %d" % step)
  
              raw_test_feats = raw_feats[:i_dev]
              raw_dev_feats = raw_feats[i_dev:i_train]
              raw_train_feats = raw_feats[i_train:]
  
              raw_test_labels = raw_labels[:i_dev]
              raw_dev_labels = raw_labels[i_dev:i_train]
              raw_train_labels = raw_labels[i_train:]
              
              run_model(
                  raw_train_feats, raw_train_labels, raw_dev_feats, raw_dev_labels, raw_test_feats, raw_test_labels,
              )
              # cycle
              raw_feats = raw_feats[i_dev:] + raw_feats[:i_dev]
              raw_labels = raw_labels[i_dev:] + raw_labels[:i_dev]
      else:
  
          print('loading training data')
          raw_train_feats, raw_train_labels = corpus_feats_and_labels(os.path.join(corpus_path, 'train'), label_scheme=scheme)
          print('loading dev data')
          raw_dev_feats, raw_dev_labels = corpus_feats_and_labels(os.path.join(corpus_path, 'dev'), label_scheme=scheme)
          print('loading test data')
          raw_test_feats, raw_test_labels = corpus_feats_and_labels(os.path.join(corpus_path, 'test'), label_scheme=scheme)
  
          (train_feats, dev_feats, test_feats), feat_indices = encode_feats(raw_train_feats, raw_dev_feats, raw_test_feats)
  
          n_feats = len(feat_indices)
          train_labels, dev_labels, test_labels = encode_labels(raw_train_labels, raw_dev_labels, raw_test_labels, scheme=scheme)
  
          model = QuoteDetectionModel(n_feats, use_crf=use_crf, sample_steps=1, label_scheme=scheme, viterbi=False, transformer=False, pipeline=False)
  
          if embedding_path is not None:
              inject_pretrained_embeddings(model, embedding_path, feat_indices)
  
          optimizer = optim.Adam(model.parameters())
          best_f1 = -1
          def training_callback():
              global best_f1
              if check_presence:
                  ev = get_ev(model, dev_feats, raw_dev_labels, eval_mode='presence')
              else:
                  ev = get_ev(model, dev_feats, raw_dev_labels)
              print(report(ev, roles=roles))
              f1 = 0
              if 'content' in ev:
                  tp = ev['content']['tp']
                  fp = ev['content']['fp']
                  fn = ev['content']['fn']
              else:
                  tp = 0
                  fp = 0
                  fn = 0
                  for role in ev:
                      tp += ev[role]['tp']
                      fp += ev[role]['fp']
                      fn += ev[role]['fn']
  
              if tp != 0:
                  p = tp / (tp + fp)
                  r = tp / (tp + fn)
                  f1 = 2 / (1/p + 1/r)
              if f1 > best_f1:
                  print('Best model so far! Saving...')
                  best_f1 = f1
                  torch.save(model.state_dict(), model_path)
              return f1
  
          train_loop(model, optimizer, train_feats, train_labels, dev_feats, dev_labels, callback=training_callback)
          print("Loading best model...")
          model.load_state_dict(torch.load(model_path))
          print("Evaluating on test-set...")
          ev = get_ev(model, test_feats, raw_test_labels)
          print(report(ev, roles=roles))    
  