class_docstr: Combines all coref modules together to find coreferent spans.
  
  Attributes:
      config (coref.config.Config): the model's configuration,
          see config.toml for the details
      epochs_trained (int): number of epochs the model has been trained for
      trainable (Dict[str, torch.nn.Module]): trainable submodules with their
          names used as keys
      training (bool): used to toggle train/eval modes
  
  Submodules (in the order of their usage in the pipeline):
      tokenizer (transformers.AutoTokenizer)
      bert (transformers.AutoModel)
      we (WordEncoder)
      rough_scorer (RoughScorer)
      pw (PairwiseEncoder)
      a_scorer (AnaphoricityScorer)
      sp (SpanPredictor)
class_docstr_tok: ["Combines","all","coref","modules","together","to","find","coreferent","spans",".","Attributes",":","config","(","coref.config.Config",")",":","the","model","'s","configuration",",","see","config.toml","for","the","details","epochs_trained","(","int",")",":","number","of","epochs","the","model","has","been","trained","for","trainable","(","Dict","[","str",",","torch.nn.Module","]",")",":","trainable","submodules","with","their","names","used","as","keys","training","(","bool",")",":","used","to","toggle","train\/eval","modes","Submodules","(","in","the","order","of","their","usage","in","the","pipeline",")",":","tokenizer","(","transformers.AutoTokenizer",")","bert","(","transformers.AutoModel",")","we","(","WordEncoder",")","rough_scorer","(","RoughScorer",")","pw","(","PairwiseEncoder",")","a_scorer","(","AnaphoricityScorer",")","sp","(","SpanPredictor",")"]
docstr: Trains all the trainable blocks in the model using the config provided.
docstr_tok: ["Trains","all","the","trainable","blocks","in","the","model","using","the","config","provided","."]
code: def train(self):
      """
      Trains all the trainable blocks in the model using the config provided.
      """
      docs = list(self._get_docs(self.config.train_data))
      docs_ids = list(range(len(docs)))
      avg_spans = sum(len(doc["head2span"]) for doc in docs) / len(docs)
  
      best_f1 = None
      for epoch in range(self.epochs_trained, self.config.train_epochs):
          self.training = True
          if self.config.log_norms:
              self.log_norms()
          running_c_loss = 0.0
          running_s_loss = 0.0
          random.shuffle(docs_ids)
          pbar = tqdm(docs_ids, unit="docs", ncols=0)
          for doc_id in pbar:
              doc = docs[doc_id]
  
              for optim in self.optimizers.values():
                  optim.zero_grad()
  
              res = self.run(doc)
  
              c_loss = self._coref_criterion(
                  res.coref_scores, res.coref_y
              )
              if res.span_y:
                  s_loss = (
                      (
                          self._span_criterion(
                              res.span_scores[:, :, 0], res.span_y[0]
                          )
                          + self._span_criterion(
                              res.span_scores[:, :, 1], res.span_y[1]
                          )
                      )
                      / avg_spans
                      / 2
                  )
              else:
                  s_loss = torch.zeros_like(c_loss)
  
              del res
  
              (c_loss + s_loss).backward()
              running_c_loss += c_loss.item()
              running_s_loss += s_loss.item()
  
              del c_loss, s_loss
  
              for optim in self.optimizers.values():
                  optim.step()
              for scheduler in self.schedulers.values():
                  scheduler.step()
  
              pbar.set_description(
                  f"Epoch {epoch + 1}:"
                  f" {doc['document_id']:26}"
                  f" c_loss: {running_c_loss / (pbar.n + 1):<.5f}"
                  f" s_loss: {running_s_loss / (pbar.n + 1):<.5f}"
              )
  
          self.epochs_trained += 1
          scores = self.evaluate()
          prev_best_f1 = best_f1
          if best_f1 is None or scores[1] > best_f1:
              if best_f1 is None:
                  logger.info(
                      "Saving new best model: F1 %.4f", scores[1]
                  )
              else:
                  logger.info(
                      "Saving new best model: F1 %.4f > %.4f",
                      scores[1],
                      best_f1,
                  )
              best_f1 = scores[1]
              # TODO: choose a different default save dir
              save_path = os.path.join(
                  self.config.save_dir, f"{self.config.section}.pt"
              )
              self.save_weights(save_path, save_optimizers=False)
          if self.config.save_each_checkpoint:
              self.save_weights()
          else:
              checkpoint_path = os.path.join(
                  self.config.save_dir,
                  f"{self.config.section}.checkpoint.pt",
              )
              self.save_weights(checkpoint_path)
          if prev_best_f1 is not None and prev_best_f1 != best_f1:
              logger.info(
                  "Epoch %d finished.\nSentence F1 %.5f p %.5f r %.5f\nBest F1 %.5f\nPrevious best F1 %.5f",
                  self.epochs_trained,
                  scores[1],
                  scores[2],
                  scores[3],
                  best_f1,
                  prev_best_f1,
              )
          else:
              logger.info(
                  "Epoch %d finished.\nSentence F1 %.5f p %.5f r %.5f\nBest F1 %.5f",
                  self.epochs_trained,
                  scores[1],
                  scores[2],
                  scores[3],
                  best_f1,
              )
  
imports: from datetime import datetime
  import json
  import logging
  import os
  import random
  import re
  from typing import Any, Dict, List, Optional, Set, Tuple
  import numpy
  import toml
  import torch
  import transformers
  from stanza.models.coref import bert, conll, utils
  from stanza.models.coref.anaphoricity_scorer import AnaphoricityScorer
  from stanza.models.coref.cluster_checker import ClusterChecker
  from stanza.models.coref.config import Config
  from stanza.models.coref.const import CorefResult, Doc
  from stanza.models.coref.loss import CorefLoss
  from stanza.models.coref.pairwise_encoder import PairwiseEncoder
  from stanza.models.coref.rough_scorer import RoughScorer
  from stanza.models.coref.span_predictor import SpanPredictor
  from stanza.models.coref.tokenizer_customization import TOKENIZER_FILTERS, TOKENIZER_MAPS
  from stanza.models.coref.utils import GraphNode
  from stanza.models.coref.word_encoder import WordEncoder
  from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
  from stanza.utils.get_tqdm import get_tqdm
label_desc: Use image node_id 12 for calling the CorefModel obj's underlying member method code with example usage: obj.train() without return types