Update app.py

app.py

@@ -100,13 +100,22 @@ examples = [
     ["filetype:pdf intitle:python"]
 ]
 
-
 class CustomLLMSettings(BaseModel):
     structured_output: LlmStructuredOutputSettings
     temperature: float = Field(default=0.7)
     top_p: float = Field(default=0.95)
     repetition_penalty: float = Field(default=1.1)
-    top_k: int = Field(default=50)  # Added top_k parameter
+    top_k: int = Field(default=50)
+    max_tokens: int = Field(default=1000)
+    stop: list[str] = Field(default_factory=list)
+    echo: bool = Field(default=False)
+    stream: bool = Field(default=False)
+    logprobs: int = Field(default=None)
+    presence_penalty: float = Field(default=0.0)
+    frequency_penalty: float = Field(default=0.0)
+    best_of: int = Field(default=1)
+    logit_bias: dict = Field(default_factory=dict)
+    max_tokens_per_summary: int = Field(default=2048)
 
 class HuggingFaceHubWrapper:
     def __init__(self, repo_id, model_kwargs, huggingfacehub_api_token):
@@ -118,7 +127,9 @@ class HuggingFaceHubWrapper:
         self.temperature = model_kwargs.get('temperature', 0.7)
         self.top_p = model_kwargs.get('top_p', 0.95)
         self.repetition_penalty = model_kwargs.get('repetition_penalty', 1.1)
-        self.top_k = model_kwargs.get('top_k', 50)  # Added top_k
+        self.top_k = model_kwargs.get('top_k', 50)
+        self.max_tokens = model_kwargs.get('max_length', 1000)
+        self.max_tokens_per_summary = model_kwargs.get('max_tokens_per_summary', 2048)
 
     def get_provider_default_settings(self):
         return CustomLLMSettings(
@@ -131,7 +142,9 @@ class HuggingFaceHubWrapper:
             temperature=self.temperature,
             top_p=self.top_p,
             repetition_penalty=self.repetition_penalty,
-            top_k=self.top_k  # Added top_k
+            top_k=self.top_k,
+            max_tokens=self.max_tokens,
+            max_tokens_per_summary=self.max_tokens_per_summary
         )
 
     def get_provider_identifier(self):
@@ -172,7 +185,7 @@ class CitingSources(BaseModel):
     )
 
 # Model function
-def get_model(temperature, top_p, repetition_penalty, top_k=50):
+def get_model(temperature, top_p, repetition_penalty, top_k=50, max_tokens=1000, max_tokens_per_summary=2048):
     return HuggingFaceHubWrapper(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
         model_kwargs={
@@ -180,7 +193,8 @@ def get_model(temperature, top_p, repetition_penalty, top_k=50):
             "top_p": top_p,
             "repetition_penalty": repetition_penalty,
             "top_k": top_k,
-            "max_length": 1000
+            "max_length": max_tokens,
+            "max_tokens_per_summary": max_tokens_per_summary
         },
         huggingfacehub_api_token=huggingface_token
     )
@@ -207,10 +221,10 @@ def respond(
     temperature,
     top_p,
     repeat_penalty,
-    top_k=50,  # Added top_k parameter
+    top_k=50,
+    max_tokens_per_summary=2048
 ):
-    model = get_model(temperature, top_p, repeat_penalty, top_k)
-    
+    model = get_model(temperature, top_p, repeat_penalty, top_k, max_tokens, max_tokens_per_summary)
     chat_template = MessagesFormatterType.MISTRAL
     
     search_tool = WebSearchTool(
@@ -262,7 +276,8 @@ demo = gr.ChatInterface(
         gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
         gr.Slider(minimum=0.0, maximum=2.0, value=1.1, step=0.1, label="Repetition penalty"),
-        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k"),  # Added top_k slider
+        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k"),
+        gr.Slider(minimum=1, maximum=4096, value=2048, step=1, label="Max tokens per summary"),
     ],
     theme=gr.themes.Soft(
         primary_hue="orange",