class_docstr: Model wrapper for coreference resolution.
class_docstr_tok: ["Model","wrapper","for","coreference","resolution","."]
code: def __init__(self, config, device):
      super().__init__()
      self.config = config
      self.max_seg_len = config["max_segment_len"]
      self.device = device  # main device
  
      # Model
      self.dropout = nn.Dropout(p=config["dropout_rate"])
  
      if "device_map" in config:
          self.config["device_map"] = {
              int(k): v for k, v in config["device_map"].items()
          }
  
      self.tz = T5Tokenizer.from_pretrained("t5-small")
  
      self.SPEAKER_START = "<speaker>"
      self.SPEAKER_END = "</speaker>"
      self.MENTION_START = "<m>"
      self.MENTION_END = "</m>"
  
      self.tz.add_tokens(self.SPEAKER_START)
      self.tz.add_tokens(self.SPEAKER_END)
      self.tz.add_tokens(self.MENTION_START)
      self.tz.add_tokens(self.MENTION_END)
  
      self.mention_start_id = self.tz.convert_tokens_to_ids(
          self.MENTION_START
      )
      self.mention_end_id = self.tz.convert_tokens_to_ids(
          self.MENTION_END
      )
  
      self.model = T5Coref.from_pretrained(
          config["plm_pretrained_name_or_path"],
          asp_hidden_dim=config["hidden_size"],
          asp_dropout_rate=config["dropout_rate"],
          asp_init_std=config["init_std"],
          asp_feature_emb_size=config["feature_emb_size"],
          asp_activation=config["activation"],
          mention_start_id=self.mention_start_id,
          mention_end_id=self.mention_end_id,
      )
  
      self.beam_size = config["beam_size"]
      self.model.resize_token_embeddings(self.tz.vocab_size + 4)
  
file_path: DeepKE/src/deepke/triple_extraction/ASP/models/model_coref.py
imports: import sys
  import os
  import torch
  import torch.nn
  import logging
  from transformers import T5Tokenizer
  from .t5_coref import T5Coref
label_desc: Use image node_id 1 to create a new CorefModel object from inherited base classes: torch.nn with example: obj = CorefModel(config, device)