Update app.py

app.py

@@ -19,10 +19,14 @@ class StoppingCriteriaSub(StoppingCriteria):
         self.encounters = encounters
 
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
+        last_tkn = input_ids[0][-1]
+        stop_word_found = False
         for stop in self.stops:
-            if sum( input_ids[0] == stop ) >= self.encounters: return True
-        return False
+            if sum( input_ids[0] == stop ) >= self.encounters: 
+                stop_word_found = True
+        return stop_word_found and self.stops[0] == last_tkn
 
+# The StoppingCriteriaSub assumes period is the first token id.
 stop_words = ['.']
 stop_words_ids = [tokenizer(stop_word, return_tensors='pt', add_special_tokens=False)['input_ids'].squeeze() for stop_word in stop_words]
 stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids,