Gagan Bhatia commited on
Commit
e4c6ec6
·
1 Parent(s): b9ee162

Update model.py

Browse files
Files changed (1) hide show
  1. src/models/model.py +1 -0
src/models/model.py CHANGED
@@ -265,6 +265,7 @@ class Summarization:
265
  self.model = MT5ForConditionalGeneration.from_pretrained(
266
  f"{model_name}", return_dict=True
267
  )
 
268
  self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}")
269
 
270
  def train(
 
265
  self.model = MT5ForConditionalGeneration.from_pretrained(
266
  f"{model_name}", return_dict=True
267
  )
268
+ elif model_type == "byt5":
269
  self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}")
270
 
271
  def train(