Fix banned tokens not working for Erato

8775247942cf86f48c1664b8bc0b119466d06d06

Cohee <18619528+Cohee1207@users.noreply.github.com>

Parents: 182756fcb2ceb37dff2d4ec615c82a76155f75e8
1 files changed, +3 -3Ignore whitespace
src/endpoints/novelai.js+3 -3
@@ -116,10 +116,10 @@ router.post('/generate', jsonParser, async function (req, res) {
116 controller.abort();116 controller.abort();
117 });117 });
118118
119 const isNewModel = (req.body.model.includes('clio') || req.body.model.includes('kayra'));119 const isNewModel = (req.body.model.includes('clio') || req.body.model.includes('kayra') || req.body.model.includes('erato'));
120 const badWordsList = getBadWordsList(req.body.model);120 const badWordsList = getBadWordsList(req.body.model);
121121
122 // Add customized bad words for Clio and Kayra122 // Add customized bad words for Clio, Kayra, and Erato
123 if (isNewModel && Array.isArray(req.body.bad_words_ids)) {123 if (isNewModel && Array.isArray(req.body.bad_words_ids)) {
124 for (const badWord of req.body.bad_words_ids) {124 for (const badWord of req.body.bad_words_ids) {
125 if (Array.isArray(badWord) && badWord.every(x => Number.isInteger(x))) {125 if (Array.isArray(badWord) && badWord.every(x => Number.isInteger(x))) {
@@ -183,7 +183,7 @@ router.post('/generate', jsonParser, async function (req, res) {
183 };183 };
184184
185 // Tells the model to stop generation at '>'185 // Tells the model to stop generation at '>'
186 if ('theme_textadventure' === req.body.prefix && isNewModel) {186 if ('theme_textadventure' === req.body.prefix && isNewModel && !req.body.model.includes('erato')) {
187 data.parameters.eos_token_id = 49405;187 data.parameters.eos_token_id = 49405;
188 }188 }
189189