updated configs

app.py

@@ -9,12 +9,17 @@ def clear():
     return None, None, None
 
 
-def generate_text(query_text, model_name: Optional[str] = "google/gemma-2b-it"):
+def generate_text(query_text, model_name: Optional[str] = "google/gemma-2b-it", tokens: Optional[int] = 1024,
+                  temp: Optional[float] = 0.49):
     combined_information = get_context(query_text)
+    gen_config = {
+        "temperature": temp,
+        "max_output_tokens": tokens,
+    }
     if model_name is None or model_name == "google/gemma-2b-it":
-        return gemma_predict(combined_information, model_name)
+        return gemma_predict(combined_information, model_name, gen_config)
     if model_name == "gemini-1.0-pro":
-        return get_gemini_response(combined_information, model_name, None)
+        return get_gemini_response(combined_information, model_name, gen_config)
     return "Sorry, something went wrong! Please try again."
 
 
@@ -42,6 +47,13 @@ with gr.Blocks() as demo:
         )
         output = gr.Textbox(label="Generated Results", lines=4)
 
+        with gr.Accordion("Settings", open=False):
+            max_new_tokens = gr.Slider(label="Max new tokens", value=1024, minimum=0, maximum=8192, step=64,
+                                       interactive=True,
+                                       visible=True, info="The maximum number of output tokens")
+            temperature = gr.Slider(label="Temperature", step=0.01, minimum=0.01, maximum=1.0, value=0.49,
+                                    interactive=True,
+                                    visible=True, info="The value used to module the logits distribution")
     with gr.Group():
         with gr.Row():
             submit_btn = gr.Button("Submit", variant="primary")
@@ -54,7 +66,7 @@ with gr.Blocks() as demo:
     gr.Markdown("## Examples")
     gr.Examples(
         examples, inputs=[query, model], label="Examples", fn=generate_text, outputs=[output],
-        cache_examples=True,
+        # cache_examples=True,
     )
 
 if __name__ == "__main__":

models/gemini.py

@@ -7,18 +7,30 @@ from dotenv import load_dotenv
 
 sys.path.append("../")
 from setup.vertex_ai_setup import initialize_vertexai_params
-from vertexai.generative_models import GenerativeModel
+from vertexai import generative_models
 
 load_dotenv()
 VERTEXAI_PROJECT = os.environ["VERTEXAI_PROJECT"]
 
+DEFAULT_GEN_CONFIG = {
+    "temperature": 0.49,
+    "max_output_tokens": 1024,
+}
+
+DEFAULT_SAFETY_SETTINGS = {
+    generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+    generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+}
 
-def get_gemini_response(prompt_text, model, parameters: Optional = None, location: Optional[str] = "us-central1") -> str:
-    initialize_vertexai_params()
 
-    if model is None or parameters is None:
+def get_gemini_response(prompt_text, model, generation_config: Optional[dict] = None,
+                        safety_settings: Optional[dict] = None) -> str:
+    initialize_vertexai_params()
+    if model is None:
         model = "gemini-1.0-pro"
-    model = GenerativeModel(model)
+    model = generative_models.GenerativeModel(model,
+                                              generation_config=DEFAULT_GEN_CONFIG if generation_config is None else generation_config,
+                                              safety_settings=DEFAULT_SAFETY_SETTINGS if safety_settings is None else safety_settings)
 
     model_response = model.generate_content(prompt_text)
 

models/gemma.py

@@ -1,4 +1,5 @@
 import os
+from typing import Optional
 
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient
@@ -6,10 +7,12 @@ from huggingface_hub import InferenceClient
 load_dotenv()
 
 
-def gemma_predict(combined_information, model_name):
+def gemma_predict(combined_information, model_name, config: Optional[dict]):
     HF_token = os.environ["HF_TOKEN"]
     client = InferenceClient(model_name, token=HF_token)
-    stream = client.text_generation(prompt=combined_information, details=True, stream=True, max_new_tokens=2048,
+    stream = client.text_generation(prompt=combined_information, details=True, stream=True,
+                                    max_new_tokens=config["max_output_tokens"],
+                                    temperature=config["temperature"],
                                     return_full_text=False)
     output = ""