Update app.py

app.py

@@ -7,26 +7,32 @@ def load_model(model_name):
 def generate(
     model_name,
     system_input,
+    user_initial_message,
+    assistant_initial_message,
     user_input,
-    temperature=0.4,
-    top_p=0.25,
-    top_k=7,
-    repetition_penalty=1.0,
-    max_new_tokens=256,
 ):
     pipe = load_model(model_name)
+    
     message_template = [
-        {
-            "role": "system",
-            "content": system_input,
-        },
-        {"role": "user", "content": "Hi!"},
-        {"role": "assistant", "content": "Hello there. How can I assist you today?"},
+        {"role": "system", "content": system_input},
+        {"role": "user", "content": user_initial_message},
+        {"role": "assistant", "content": assistant_initial_message},
         {"role": "user", "content": user_input},
     ]
+    
     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=repetition_penalty)
+    
+    if model_name == "Felladrin/Pythia-31M-Chat-v1":
+        outputs = pipe(prompt, max_new_tokens=1024, do_sample=True, temperature=0.4, top_k=7, top_p=0.25, repetition_penalty=1.0016)
+    elif model_name == "Felladrin/Llama-160M-Chat-v1":
+        outputs = pipe(prompt, max_new_tokens=1024, penalty_alpha=0.5, top_k=5, repetition_penalty=1.01)
+    elif model_name == "Felladrin/TinyMistral-248M-SFT-v4":
+        outputs = pipe(prompt, max_new_tokens=1024, penalty_alpha=0.5, top_k=5, repetition_penalty=1.0)
+    elif model_name == "Felladrin/Smol-Llama-101M-Chat-v1":
+        outputs = pipe(prompt, max_new_tokens=1024, penalty_alpha=0.5, top_k=5, repetition_penalty=1.0, add_special_tokens=True)
+    else:
+        outputs = pipe(prompt, max_new_tokens=1024, do_sample=True, temperature=0.9, top_k=50, top_p=0.95, repetition_penalty=1.2)
+        
     return outputs[0]["generated_text"]
 
 model_choices = ["Felladrin/Llama-160M-Chat-v1", "Felladrin/Smol-Llama-101M-Chat-v1", "Felladrin/TinyMistral-248M-SFT-v4", "Felladrin/Pythia-31M-Chat-v1"]
@@ -36,12 +42,9 @@ g = gr.Interface(
     inputs=[
         gr.components.Dropdown(choices=model_choices, label="Model", value=model_choices[0], interactive=True),
         gr.components.Textbox(lines=2, label="System Message", value="You are a highly knowledgeable and friendly assistant. Your goal is to understand and respond to user inquiries with accuracy and clarity. You're adept at providing detailed explanations, concise summaries, and insightful responses. Your interactions are always respectful, helpful, and focused on delivering the most relevant information to the user."),
+        gr.components.Textbox(lines=2, label="User Initial Message", value="Hi!"),
+        gr.components.Textbox(lines=2, label="Assistant Initial Message", value="Hello there. How can I assist you today?"),
         gr.components.Textbox(lines=2, label="User Message", value="Tell me something curious about the Earth!"),
-        gr.components.Slider(minimum=0, maximum=1, value=0.4, label="Temperature"),
-        gr.components.Slider(minimum=0, maximum=1, value=0.25, label="Top p"),
-        gr.components.Slider(minimum=0, maximum=50, step=1, value=7, label="Top k"),
-        gr.components.Slider(minimum=1.0, maximum=1.5, step=0.001, value=1.0016, label="Repetition Penalty"),
-        gr.components.Slider(minimum=1, maximum=1024, step=1, value=256, label="Max tokens"),
     ],
     outputs=[gr.Textbox(lines=10, label="Output")],
     title="A place to try out text-generation models fine-tuned by Felladrin",