Add prefix and source file options

README.md

@@ -1,5 +1,5 @@
 ---
-title: Prompt Generator
+title: Test Prompt Generator
 emoji: ⚡
 colorFrom: blue
 colorTo: green

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from prompt_generator import generate_prompt
+from test_prompt_generator import generate_prompt
 
 tokenizers = [
     "google/gemma-7b",
@@ -16,20 +16,36 @@ tokenizers = [
 ]
 
 
-def generate(model_id, num_tokens):
+def generate(model_id, num_tokens, prefix=None, source_text=None):
     output_file = f"prompt_{num_tokens}.jsonl"
-    prompt = generate_prompt(model_id, int(num_tokens), silent=True, output_file=output_file)
+    prompt = generate_prompt(model_id, int(num_tokens), prefix=prefix, source_text=source_text, output_file=output_file)
     return prompt, output_file
 
 
 demo = gr.Interface(
     fn=generate,
     title="Test Prompt Generator",
-    description="Generate prompts with a given length for testing transformer models. "
+    description="Generate prompts with a given number of tokens for testing transformer models. "
     "Prompt source: https://archive.org/stream/alicesadventures19033gut/19033.txt",
     inputs=[
-        gr.Dropdown(label="Tokenizer", choices=tokenizers, allow_custom_value=True),
-        gr.Textbox(label="Number of Tokens"),
+        gr.Dropdown(
+            label="Tokenizer",
+            choices=tokenizers,
+            value="mistralai/Mistral-7B-v0.1",
+            allow_custom_value=True,
+            info="Select a tokenizer from this list or paste a model_id from a model on the Hugging Face Hub",
+        ),
+        gr.Number(
+            label="Number of Tokens", minimum=4, maximum=2048, value=32, info="Enter a number between 4 and 2048."
+        ),
+        gr.Textbox(
+            label="Prefix (optional)",
+            info="If given, the start of the prompt will be this prefix. Example: 'Summarize the following text:'",
+        ),
+        gr.Textbox(
+            label="Source text (optional)",
+            info="By default, prompts will be generated from Alice in Wonderland. Enter text here to use that instead.",
+        ),
     ],
     outputs=[gr.Textbox(label="prompt", show_copy_button=True), gr.File(label="Json file")],
     examples=[

requirements.txt

@@ -1,2 +1,3 @@
 transformers
-git+https://github.com/helena-intel/prompt-generator.git
+sentencepiece
+git+https://github.com/helena-intel/test-prompt-generator.git