Update app.py

app.py

@@ -30,7 +30,10 @@ def generation(tokenizer, model, content):
     typical_p = 0.8 if decoder == 'Typical' else None
     do_sample = (decoder in ['Beam', 'Typical', 'Sample'])
 
-    seek_token_ids = tokenizer.encode('= " +')[1:]
+    seek_token_ids = [
+        tokenizer.encode('= \'" +')[1:],
+        tokenizer.encode('= " +')[1:],
+    ]
 
     full_output = model.generate(content, generate=6, do_sample=False)
 
@@ -49,7 +52,9 @@ def generation(tokenizer, model, content):
             return my_token_prob * next_words(newprompt, position + 1, seek_token_ids[1:])
         return my_token_prob
 
-    prob = next_words(content, len(tokenizer(content)['input_ids']), seek_token_ids)
+    prob = 0
+    for opt in seek_token_ids:
+        prob += next_words(content, len(tokenizer(content)['input_ids']), opt)
     return ["".join(full_output.tokens), str(prob.item() * 100) + '% chance of risky concatenation']
 
 def code_from_prompts(prompt, model, type_hints, pre_content):