add context

app.py

@@ -13,7 +13,7 @@ model = AutoModelForCausalLM.from_pretrained(model_path)
 pipe = pipeline("text-generation", model=model_path)
 pipe.tokenizer = tokenizer  # Assign tokenizer manually
 
-def create_prompt(system_message, user_message, tool_definition=""):
+def create_prompt(system_message, user_message, tool_definition="", context=""):
     if tool_definition:
         return f"""<extra_id_0>System
 {system_message}
@@ -22,7 +22,7 @@ def create_prompt(system_message, user_message, tool_definition=""):
 {tool_definition}
 </tool>
 <context>
-The current date is 2023-06-01.
+{context}
 </context>
 
 <extra_id_1>User
@@ -33,8 +33,8 @@ The current date is 2023-06-01.
         return f"<extra_id_0>System\n{system_message}\n\n<extra_id_1>User\n{user_message}\n<extra_id_1>Assistant\n"
 
 @spaces.GPU
-def generate_response(message, history, system_message, max_tokens, temperature, top_p, use_pipeline=False, tool_definition=""):
-    full_prompt = create_prompt(system_message, message, tool_definition)
+def generate_response(message, history, system_message, max_tokens, temperature, top_p, use_pipeline=False, tool_definition="", context=""):
+    full_prompt = create_prompt(system_message, message, tool_definition, context)
     
     if use_pipeline:
         messages = [
@@ -73,22 +73,28 @@ def generate_response(message, history, system_message, max_tokens, temperature,
     return assistant_response
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Nemotron-Mini-4B-Instruct Demo with Custom Function Calling")
-    gr.Markdown("This demo showcases the Nemotron-Mini-4B-Instruct model from NVIDIA, including optional custom function calling.")
-    
+    with gr.Row(): 
+        gr.Markdown(title)
+    with gr.Row():
+        gr.Markdown(description)    
     with gr.Row():
         with gr.Column(scale=3):
             chatbot = gr.Chatbot(height=400)
             msg = gr.Textbox(label="User Input", placeholder="Ask a question or request a task...")
-            clear = gr.Button("Clear")
-        
-        with gr.Column(scale=2):
+            with gr.Row():
+                send = gr.Button("Send")
+                clear = gr.Button("Clear")
             system_message = gr.Textbox(
                 label="System Message",
                 value="You are a helpful AI assistant.",
                 lines=2,
                 placeholder="Set the AI's behavior and context..."
             )
+            context = gr.Textbox(
+                label="Context",
+                lines=2,
+                placeholder="Enter additional context information..."
+            )
             max_tokens = gr.Slider(minimum=1, maximum=1024, value=256, step=1, label="Max Tokens")
             temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
             top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
@@ -100,20 +106,22 @@ with gr.Blocks() as demo:
                     value=customtool,
                     lines=15,
                     language="json"
-                    # placeholder="Enter the JSON definition of your custom tool..."
                 )
 
     def user(user_message, history):
         return "", history + [[user_message, None]]
 
-    def bot(history, system_message, max_tokens, temperature, top_p, use_pipeline, tool_definition):
+    def bot(history, system_message, max_tokens, temperature, top_p, use_pipeline, tool_definition, context):
         user_message = history[-1][0]
-        bot_message = generate_response(user_message, history, system_message, max_tokens, temperature, top_p, use_pipeline, tool_definition)
+        bot_message = generate_response(user_message, history, system_message, max_tokens, temperature, top_p, use_pipeline, tool_definition, context)
         history[-1][1] = bot_message
         return history
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, [chatbot, system_message, max_tokens, temperature, top_p, use_pipeline, tool_definition], chatbot
+        bot, [chatbot, system_message, max_tokens, temperature, top_p, use_pipeline, tool_definition, context], chatbot
+    )
+    send.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, system_message, max_tokens, temperature, top_p, use_pipeline, tool_definition, context], chatbot
     )
     clear.click(lambda: None, None, chatbot, queue=False)
 

globe.py

@@ -39,4 +39,23 @@ customtool = """{
     },
     "required": ["param1"]
   }
-}"""
+}"""
+
+example = """{{
+  "name": "get_current_weather",
+  "description": "Get the current weather in a given location",
+  "parameters": {{
+    "type": "object",
+    "properties": {{
+      "location": {{
+        "type": "string",
+        "description": "The city and state, e.g. San Francisco, CA"
+      }},
+      "unit": {{
+        "type": "string",
+        "enum": ["celsius", "fahrenheit"]
+      }}
+    }},
+    "required": ["location"]
+  }}
+}}"""