Update app.py

app.py

@@ -12,10 +12,10 @@ model = AutoModelForCausalLM.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 @spaces.GPU
-def generate(prompt, history):
+def generate(message, chat_history, temperature=0.7, top_p=0.9, top_k=50, max_new_tokens=512, repetition_penalty=1.1):
     messages = [
         {"role": "system", "content": "You are Zurich, a 7 billion parameter Large Language model built on the Qwen 2.5 7B model developed by Alibaba Cloud, and fine-tuned by Ruben Roy. You have been fine-tuned with the GammaCorpus v2 dataset, a dataset filled with structured and filtered multi-turn conversations and was also created by Ruben Roy. You are a helpful assistant."},
-        {"role": "user", "content": prompt}
+        {"role": "user", "content": message}
     ]
     text = tokenizer.apply_chat_template(
         messages,
@@ -25,7 +25,12 @@ def generate(prompt, history):
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
     generated_ids = model.generate(
         **model_inputs,
-        max_new_tokens=512
+        temperature=float(temperature),
+        top_p=float(top_p),
+        top_k=int(top_k),
+        max_new_tokens=int(max_new_tokens),
+        repetition_penalty=float(repetition_penalty),
+        do_sample=True if float(temperature) > 0 else False
     )
     generated_ids = [
         output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
@@ -78,6 +83,28 @@ TITLE_HTML = """
         font-size: 0.9rem;
         color: #94a3b8;
     }
+    .settings-section {
+        background: rgba(255, 255, 255, 0.05);
+        padding: 1.5rem;
+        border-radius: 1rem;
+        margin: 1.5rem auto;
+        border: 1px solid rgba(255, 255, 255, 0.1);
+        max-width: 800px;
+    }
+    .settings-title {
+        color: #e2e8f0;
+        font-size: 1.25rem;
+        font-weight: 600;
+        margin-bottom: 1rem;
+        display: flex;
+        align-items: center;
+        gap: 0.7rem;
+    }
+    .parameter-info {
+        color: #94a3b8;
+        font-size: 0.8rem;
+        margin-top: 0.25rem;
+    }
 </style>
 
 <div style="background: linear-gradient(135deg, #1e293b 0%, #0f172a 100%); padding: 1.5rem; border-radius: 1.5rem; text-align: center; margin: 1rem auto; max-width: 1200px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);">
@@ -94,7 +121,7 @@ TITLE_HTML = """
 
     <div style="display: flex; gap: 1.5rem; justify-content: center;">
         <div class="model-section">
-            <h2 style="font-size: 1.25rem; color: #e2e8f0; margin-bottom: 1rem; margin-top: 1px; font-weight: 600; display: flex; align-items: center; justify-content: center; gap: 0.7rem;">
+            <h2 style="font-size: 1.25rem; color: #e2e8f0; margin-bottom: 1.4rem; margin-top: 1px; font-weight: 600; display: flex; align-items: center; justify-content: center; gap: 0.7rem;">
                 <i class="fas fa-brain"></i>
                 7B Models
             </h2>
@@ -108,7 +135,7 @@ TITLE_HTML = """
             </div>
         </div>
         <div class="model-section">
-            <h2 style="font-size: 1.25rem; color: #e2e8f0; margin-bottom: 1rem; margin-top: 1px; font-weight: 600; display: flex; align-items: center; justify-content: center; gap: 0.7rem;">
+            <h2 style="font-size: 1.25rem; color: #e2e8f0; margin-bottom: 1.4rem; margin-top: 1px; font-weight: 600; display: flex; align-items: center; justify-content: center; gap: 0.7rem;">
                 <i class="fas fa-rocket"></i>
                 14B Models
             </h2>
@@ -125,9 +152,76 @@ TITLE_HTML = """
 </div>
 """
 
+examples = [
+    ["Explain quantum computing in simple terms"],
+    ["Write a short story about a time traveler"],
+    ["Explain the process of photosynthesis"],
+    ["Tell me an intersting fact about Palm trees"]
+]
+
 with gr.Blocks() as demo:
     gr.HTML(TITLE_HTML)
-    chat_interface = gr.ChatInterface(
+    
+    with gr.Accordion("Generation Settings", open=False):
+        with gr.Row():
+            with gr.Column():
+                temperature = gr.Slider(
+                    minimum=0.0,
+                    maximum=2.0,
+                    value=0.7,
+                    step=0.1,
+                    label="Temperature",
+                    info="Higher values make the output more random, lower values make it more deterministic",
+                    interactive=True
+                )
+                top_p = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.05,
+                    label="Top P",
+                    info="Controls the cumulative probability threshold for nucleus sampling",
+                    interactive=True
+                )
+                top_k = gr.Slider(
+                    minimum=1,
+                    maximum=100,
+                    value=50,
+                    step=1,
+                    label="Top K",
+                    info="Limits the number of tokens to consider for each generation step",
+                    interactive=True
+                )
+            with gr.Column():
+                max_new_tokens = gr.Slider(
+                    minimum=1,
+                    maximum=2048,
+                    value=512,
+                    step=1,
+                    label="Max New Tokens",
+                    info="Maximum number of tokens to generate in the response",
+                    interactive=True
+                )
+                repetition_penalty = gr.Slider(
+                    minimum=1.0,
+                    maximum=2.0,
+                    value=1.1,
+                    step=0.1,
+                    label="Repetition Penalty",
+                    info="Higher values stop the model from repeating the same info",
+                    interactive=True
+                )
+    
+    chatbot = gr.ChatInterface(
         fn=generate,
+        additional_inputs=[
+            temperature,
+            top_p,
+            top_k,
+            max_new_tokens,
+            repetition_penalty
+        ],
+        examples=examples
     )
+
 demo.launch(share=True)