AndriiPets commited on
Commit
406142b
1 Parent(s): f956d85

chef choice

Browse files
Files changed (1) hide show
  1. generator.py +3 -2
generator.py CHANGED
@@ -90,9 +90,10 @@ def skip_special_tokens_and_prettify(text):
90
  return data
91
 
92
 
93
- def generation_function(texts, lang="en"):
94
 
95
  langs = ['ru', 'en']
 
96
  api_credentials = load_api()
97
 
98
  if lang != "en" and lang in langs:
@@ -100,7 +101,7 @@ def generation_function(texts, lang="en"):
100
 
101
  output_ids = generator(
102
  texts,
103
- ** chef_top
104
  )[0]["generated_token_ids"]
105
 
106
  recepie = tokenizer.decode(output_ids, skip_special_tokens=False)
 
90
  return data
91
 
92
 
93
+ def generation_function(texts, lang="en", chef="top"):
94
 
95
  langs = ['ru', 'en']
96
+ chefs = {'top': chef_top, 'beam': chef_beam}
97
  api_credentials = load_api()
98
 
99
  if lang != "en" and lang in langs:
 
101
 
102
  output_ids = generator(
103
  texts,
104
+ ** chefs[chef]
105
  )[0]["generated_token_ids"]
106
 
107
  recepie = tokenizer.decode(output_ids, skip_special_tokens=False)