Update app.py

app.py

@@ -9,19 +9,18 @@ def generate(
     template_name,
     user_input,
     temperature=0.4,
-    top_p=0.95,
-    top_k=50,
+    top_p=0.25,
+    top_k=7,
     max_new_tokens=256,
 ):
     pipe = load_model(model_name)
-    # Need to add additional options later.
-    if template_name == "Falcon 1B Template":
+    if template_name == "User-Assistant":
         message_template = [
             {"role": "user", "content": "Hello!"},
-            {"role": "assistant", "content": "Hello! How can I assist you today?"},
+            {"role": "assistant", "content": "Hi! How can I assist you today?"},
             {"role": "user", "content": user_input},
         ]
-    else:  # Default to "TinyLlama Template"
+    else:
         message_template = [
             {
                 "role": "system",
@@ -30,15 +29,14 @@ def generate(
             {"role": "user", "content": user_input},
         ]
 
-    # Set tokenize correctly. Otherwise ticking the box breaks it.
     prompt = pipe.tokenizer.apply_chat_template(message_template, tokenize=False, add_generation_prompt=True)
     outputs = pipe(prompt, max_new_tokens=max_new_tokens, do_sample=True,
                    temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=1.10)
     return outputs[0]["generated_text"]
 
-model_choices = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0", "ericzzz/falcon-rw-1b-chat"]
-template_choices = ["TinyLlama Template", "Falcon Template"]
-# What at the best options? 
+model_choices = ["Felladrin/Pythia-31M-Chat-v1", "Felladrin/Llama-160M-Chat-v1", "Felladrin/Smol-Llama-101M-Chat-v1", "Felladrin/TinyMistral-248M-SFT-v4"]
+template_choices = ["System-User-Assistant", "User-Assistant"]
+
 g = gr.Interface(
     fn=generate,
     inputs=[
@@ -46,13 +44,13 @@ g = gr.Interface(
         gr.components.Dropdown(choices=template_choices, label="Template", value=template_choices[0], interactive=True),
         gr.components.Textbox(lines=2, label="Prompt", value="How many planets are in our solar system?"),
         gr.components.Slider(minimum=0, maximum=1, value=0.4, label="Temperature"),
-        gr.components.Slider(minimum=0, maximum=1, value=0.95, label="Top p"),
-        gr.components.Slider(minimum=0, maximum=100, step=1, value=50, label="Top k"),
+        gr.components.Slider(minimum=0, maximum=1, value=0.25, label="Top p"),
+        gr.components.Slider(minimum=0, maximum=100, step=1, value=7, label="Top k"),
         gr.components.Slider(minimum=1, maximum=1024, step=1, value=256, label="Max tokens"),  
     ],
     outputs=[gr.Textbox(lines=10, label="Output")],
-    title="Hugging Face Transformers Model",
-    description="A simple interface for generating text with a Hugging Face Transformers model.",
+    title="Chat with Felladrin's LLMs",
+    description="Note that the inference happens on free-tier hardware, which may lead to slower outputs during periods of high demand.",
     concurrency_limit=1
 )