customize token output number

app.py

@@ -288,7 +288,7 @@ class AbliterationProcessor:
                 modified_weight = down_proj_weight - scale_factor * torch.matmul(projection_matrix, down_proj_weight)
                 layer.mlp.down_proj.weight.data = modified_weight
     
-    def chat(self, message, history):
+    def chat(self, message, history, max_new_tokens=2048):
         """Chat functionality"""
         if self.model is None or self.tokenizer is None:
             return "⚠️ Please load a model first!", history
@@ -335,7 +335,7 @@ class AbliterationProcessor:
             
             gen = self.model.generate(
                 toks.to(self.model.device), 
-                max_new_tokens=2048,
+                max_new_tokens=max_new_tokens,
                 temperature=0.7,
                 do_sample=True,
                 pad_token_id=self.tokenizer.eos_token_id,
@@ -472,6 +472,16 @@ def create_interface():
                         org_token.render()
                         private_repo.render()
                         
+                        gr.Markdown("### 💬 Chat Settings")
+                        max_new_tokens = gr.Number(
+                            value=2048,
+                            label="Max New Tokens",
+                            minimum=1,
+                            maximum=8192,
+                            step=1,
+                            info="Maximum number of tokens to generate in chat responses"
+                        )
+                        
                         process_btn = gr.Button("🚀 Start Processing", variant="primary")
                         process_output = gr.Markdown(label="Processing Result")
                         process_image = gr.Image(show_label=False)
@@ -538,18 +548,18 @@ def create_interface():
         def user(user_message, history):
             return "", history + [{"role": "user", "content": user_message}]
         
-        def bot(history):
+        def bot(history, max_new_tokens):
             if history and history[-1]["role"] == "user":
-                response, _ = processor.chat(history[-1]["content"], history[:-1])
+                response, _ = processor.chat(history[-1]["content"], history[:-1], max_new_tokens)
                 history.append({"role": "assistant", "content": response})
             return history
         
         msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-            bot, chatbot, chatbot
+            bot, [chatbot, max_new_tokens], chatbot
         )
         
         send_btn.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-            bot, chatbot, chatbot
+            bot, [chatbot, max_new_tokens], chatbot
         )
         
         clear.click(lambda: [], None, chatbot, queue=False)