alexkueck commited on
Commit
71ea24c
·
1 Parent(s): 2974ec1

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +11 -0
utils.py CHANGED
@@ -78,6 +78,17 @@ def load_tokenizer_and_model(base_model, load_8bit=False):
78
  model.eval()
79
  return tokenizer,model, device
80
 
 
 
 
 
 
 
 
 
 
 
 
81
  # Greedy Search
82
  def greedy_search(input_ids: torch.Tensor,
83
  model: torch.nn.Module,
 
78
  model.eval()
79
  return tokenizer,model, device
80
 
81
+
82
+ def load_tokenizer(base_model):
83
+ if torch.cuda.is_available():
84
+ device = "cuda"
85
+ else:
86
+ device = "cpu"
87
+
88
+ tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast = True)
89
+ return tokenizer
90
+
91
+
92
  # Greedy Search
93
  def greedy_search(input_ids: torch.Tensor,
94
  model: torch.nn.Module,