Update app.py

app.py

@@ -6,19 +6,7 @@ tokenizer = AutoTokenizer.from_pretrained("milyiyo/paraphraser-german-mt5-small"
 model = AutoModelForSeq2SeqLM.from_pretrained("milyiyo/paraphraser-german-mt5-small")
 
 
-def decode_model_outputs(model_outputs, model_input):
-    input_sent = model_input.split(':',1)[-1].strip()
-    sentences = []
-    for output in model_outputs:
-        sentences.append(tokenizer.decode(output, skip_special_tokens=True))
-
-    sentences = set(sentences)
-    for sent in sentences:
-        if sent != input_sent:
-            print(f'    - {sent}')
-
-
-def generate_v1(input):
+def generate_v1(input, count):
     """Generate text using a Beam Search strategy with repetition penalty."""
     model_outputs = model.generate(input["input_ids"], 
                                            early_stopping=True,
@@ -27,7 +15,8 @@ def generate_v1(input):
                                            no_repeat_ngram_size=2, 
                                            num_beams=10, 
                                            num_return_sequences=5, 
-                                           repetition_penalty=3.5, 
+                                           repetition_penalty=3.5,
+                                           num_return_sequences=count
                                            )
     sentences = []
     for output in model_outputs:
@@ -35,7 +24,7 @@ def generate_v1(input):
     return sentences
 
 
-def generate_v2(input):
+def generate_v2(input, count):
     """Generate text using a Beam Search strategy."""
     model_outputs = model.generate(input["input_ids"],
                                            early_stopping=True,
@@ -45,6 +34,7 @@ def generate_v2(input):
                                            num_beams=5,
                                            num_return_sequences=5,
                                            temperature=1.5,
+                                           num_return_sequences=count
                                           )
     sentences = []
     for output in model_outputs:
@@ -52,7 +42,7 @@ def generate_v2(input):
     return sentences
 
 
-def generate_v3(input):
+def generate_v3(input, count):
     """Generate text using a Diverse Beam Search strategy."""
     model_outputs = model.generate(input["input_ids"],
                                 num_beams=5,
@@ -63,14 +53,15 @@ def generate_v3(input):
                                 diversity_penalty=2.0,
                                 no_repeat_ngram_size=2,
                                 early_stopping=True,
-                                length_penalty=2.0)
+                                length_penalty=2.0,
+                                num_return_sequences=count)
     sentences = []
     for output in model_outputs:
         sentences.append(tokenizer.decode(output, skip_special_tokens=True))
     return sentences
 
 
-def generate_v4(input):
+def generate_v4(input, count):
     """Generate text using a Diverse Beam Search strategy."""
     input_ids, attention_masks = input["input_ids"], input["attention_mask"]
     outputs = model.generate(
@@ -80,7 +71,7 @@ def generate_v4(input):
         top_k=120,
         top_p=0.95,
         early_stopping=True,
-        num_return_sequences=p_count
+        num_return_sequences=count
     )
     res = []
     for output in outputs:
@@ -114,10 +105,10 @@ def paraphrase(sentence: str, count: str):
     #     res.append(line)
     return {
         'result': {
-            'generate_v1':generate_v1(encoding),
-            'generate_v2':generate_v2(encoding),
-            'generate_v3':generate_v3(encoding),
-            'generate_v4':generate_v4(encoding),
+            'generate_v1':generate_v1(encoding, count),
+            'generate_v2':generate_v2(encoding, count),
+            'generate_v3':generate_v3(encoding, count),
+            'generate_v4':generate_v4(encoding, count),
         }
     }