Crystalcareai commited on
Commit
3b797cc
1 Parent(s): b983c45

Update generate.py

Browse files
Files changed (1) hide show
  1. generate.py +6 -6
generate.py CHANGED
@@ -80,7 +80,7 @@ def custom_generate(
80
  # Check if the end token is generated
81
  if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("</s>"):
82
  finished_generating[answer_idx] = 1
83
-
84
  if finished_generating.all():
85
  break
86
 
@@ -89,17 +89,17 @@ def custom_generate(
89
 
90
  from collections import namedtuple
91
  GenerateOutput = namedtuple("GenerateOutput", ["sequences", "scores", "attentions", "hidden_states"])
92
-
93
  # Convert the generated token IDs to a tensor
94
  generated_token_ids_tensor = input_ids
95
-
96
  output = GenerateOutput(
97
  sequences=generated_token_ids_tensor,
98
  scores=None,
99
  attentions=None,
100
- hidden_states=None
101
  )
102
-
103
  return output
104
 
105
 
@@ -211,4 +211,4 @@ def generate(
211
  **model_kwargs,
212
  )
213
 
214
- return output
 
80
  # Check if the end token is generated
81
  if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("</s>"):
82
  finished_generating[answer_idx] = 1
83
+
84
  if finished_generating.all():
85
  break
86
 
 
89
 
90
  from collections import namedtuple
91
  GenerateOutput = namedtuple("GenerateOutput", ["sequences", "scores", "attentions", "hidden_states"])
92
+
93
  # Convert the generated token IDs to a tensor
94
  generated_token_ids_tensor = input_ids
95
+
96
  output = GenerateOutput(
97
  sequences=generated_token_ids_tensor,
98
  scores=None,
99
  attentions=None,
100
+ hidden_states=None,
101
  )
102
+
103
  return output
104
 
105
 
 
211
  **model_kwargs,
212
  )
213
 
214
+ return output