atrost commited on
Commit
307e1ed
·
1 Parent(s): d719bac

fix model type

Browse files
Files changed (1) hide show
  1. local_perplexity.py +2 -2
local_perplexity.py CHANGED
@@ -17,7 +17,7 @@ import datasets
17
  import numpy as np
18
  import torch
19
  from torch.nn import CrossEntropyLoss
20
- from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM
21
 
22
  import evaluate
23
  from evaluate import logging
@@ -107,7 +107,7 @@ class Perplexity(evaluate.Metric):
107
  else:
108
  device = "cuda" if torch.cuda.is_available() else "cpu"
109
 
110
- model = AutoModelForSeq2SeqLM.from_pretrained(model_id, local_files_only=local_file_only)
111
  model = model.to(device)
112
 
113
  tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=local_file_only)
 
17
  import numpy as np
18
  import torch
19
  from torch.nn import CrossEntropyLoss
20
+ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModel
21
 
22
  import evaluate
23
  from evaluate import logging
 
107
  else:
108
  device = "cuda" if torch.cuda.is_available() else "cpu"
109
 
110
+ model = AutoModel.from_pretrained(model_id, local_files_only=local_file_only)
111
  model = model.to(device)
112
 
113
  tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=local_file_only)