Update app.py

app.py

@@ -88,28 +88,35 @@ tokenizer.save_pretrained("./bias-eliminator-model")
 bias_prompt_eliminator = pipeline("text-generation", model="./bias-eliminator-model", tokenizer="./bias-eliminator-model")
 
 def show_neutralized_prompt(input_text):
+    # input into retrained gpt2 model requires the format:
+    # "<input_text><text sep>"
+    #
+    # Where: <input_text> is the user prompt
+    #        <text sep> is the string " -> "
+    #
+    # Example:
+    #
+    # <input text> = "Explain why immigrants struggle with career advancement in public services."
+    # Input format to model is:
+    #  <input_text><text sep> = "Explain why immigrants struggle with career advancement in public services. ->"
+
     sep = " -> "
     input_text_format = input_text + sep
-
-    result = bias_prompt_eliminator(
-        input_text_format,
-        max_length=60,
-        do_sample=False
-    )
+    result = bias_prompt_eliminator(input_text_format, max_length=30, num_return_sequences=1)
 
     generated_text = result[0]['generated_text']
 
-    if sep in generated_text:
-        output = generated_text.split(sep)[-1].strip()
+    first = generated_text.find(sep)
 
+    if first != -1:
+        second = generated_text.find(sep, first +len(sep))
+    else:
+        second = -1
+    if second != -1:
+        print(generated_text[0:second])
+    else:
+        print(generated_text[0:first])
         
-        if len(output.split()) < 3:
-            return "Rewrite using neutral, inclusive language."
-
-        return output.capitalize()
-
-    return "Could not generate neutral version."
-
 # FAIRNESS MODEL (MNLI)
 mnli_model_name = "facebookAI/roberta-large-mnli"
 mnli_tokenizer = AutoTokenizer.from_pretrained(mnli_model_name)