Update app.py

app.py

@@ -9,40 +9,41 @@ gpt2_pipe = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Dif
 with open("ideas.txt", "r") as f:
     lines = f.readlines()
 
-def generate_prompt(starting_text):
-    for count in range(4):
-        seed = random.randint(100, 1000000)
-        set_seed(seed)
-
-        # Choose a random line from the file if the input text is empty
-        if starting_text == "":
-            starting_text = lines[random.randrange(0, len(lines))].strip().lower().capitalize()
-            starting_text = re.sub(r"[,:\-–.!;?_]", '', starting_text)
-
-        # Generate text
-        response = gpt2_pipe(starting_text, max_length=random.randint(60, 90), num_return_sequences=4)
-        response_list = []
-        for x in response:
-            resp = x['generated_text'].strip()
-            if resp != starting_text and len(resp) > (len(starting_text) + 4) and not resp.endswith((":", "-", "—")):
-                response_list.append(resp + '\n')
-
-        # Clean the generated text
-        response_end = "\n".join(response_list)
-        response_end = re.sub(r'[^ ]+\.[^ ]+', '', response_end)  # Removes strings like 'abc.xyz'
-        response_end = response_end.replace("<", "").replace(">", "")
-
-        if response_end:
-            return response_end
-        if count == 4:
-            return response_end
+def generate_prompts(starting_text, num_prompts=1):
+    response_list = []
+    
+    for _ in range(num_prompts):
+        for count in range(4):  # Attempt up to 4 times to generate valid response
+            seed = random.randint(100, 1000000)
+            set_seed(seed)
+
+            # Choose a random line from the file if the input text is empty
+            if starting_text == "":
+                starting_text = lines[random.randrange(0, len(lines))].strip().lower().capitalize()
+                starting_text = re.sub(r"[,:\-–.!;?_]", '', starting_text)
+
+            # Generate text
+            response = gpt2_pipe(starting_text, max_length=random.randint(60, 90), num_return_sequences=1)
+            generated_text = response[0]['generated_text'].strip()
+
+            # Clean and check the generated response
+            if generated_text != starting_text and len(generated_text) > (len(starting_text) + 4):
+                cleaned_text = re.sub(r'[^ ]+\.[^ ]+', '', generated_text)  # Remove strings like 'abc.xyz'
+                cleaned_text = cleaned_text.replace("<", "").replace(">", "")
+                response_list.append(cleaned_text)
+                break  # Stop trying further once a valid prompt is added
+
+    return response_list[:num_prompts]
 
 # Define the API endpoint
 @app.route('/', methods=['GET'])
 def generate_api():
     starting_text = request.args.get('text', default="", type=str)
-    result = generate_prompt(starting_text)
-    return jsonify({"generated_text": result})
+    num_prompts = request.args.get('n', default=1, type=int)  # Get the number of prompts to return, default is 1
+    
+    # Generate the prompts
+    results = generate_prompts(starting_text, num_prompts=num_prompts)
+    return jsonify(results)
 
 if __name__ == '__main__':
     # Run the Flask app on port 7860