Update app.py

app.py

@@ -5,8 +5,8 @@
 #  This script has been updated to run as a Hugging Face Space.
 #
 #  Key Upgrades from the original script:
-#  1.  **Hugging Face Model Integration**: Uses the 'google/gemma-3b-it' model
-#      from the Hugging Face Hub for argument extraction, instead of the Gemini API.
+#  1.  **Hugging Face Model Integration**: Uses the 'google/gemma-3n-E4B' model
+#      from the Hugging Face Hub for argument extraction.
 #  2.  **Environment Variable Management**: Securely accesses the
 #      HUGGING_FACE_HUB_TOKEN using os.environ.get(), which is the standard
 #      for Hugging Face Spaces.
@@ -45,16 +45,18 @@ try:
         raise ValueError("HUGGING_FACE_HUB_TOKEN secret not found.")
 
     print("⚙️ Loading Hugging Face model for argument extraction...")
-    # Using a smaller, instruction-tuned model for efficient argument extraction
-    hf_tokenizer = AutoTokenizer.from_pretrained("google/gemma-3b-it", token=HF_TOKEN)
+    # Using the user-specified Gemma 3n model
+    model_id = "google/gemma-3n-E4B"
+
+    hf_tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
     hf_model = AutoModelForCausalLM.from_pretrained(
-        "google/gemma-3b-it",
+        model_id,
         token=HF_TOKEN,
         torch_dtype=torch.bfloat16, # Use bfloat16 for efficiency
         device_map="auto" # Automatically use GPU if available
     )
     USE_HF_LLM = True
-    print("✅ Successfully loaded 'google/gemma-3b-it' model.")
+    print(f"✅ Successfully loaded '{model_id}' model.")
 
 except Exception as e:
     USE_HF_LLM = False
@@ -86,7 +88,7 @@ class Tool:
         """
         schema_str = json.dumps(self.args_schema, indent=2)
         examples_str = "\n".join([f" - Example: {ex['prompt']} -> Args: {json.dumps(ex['args'])}" for ex in self.examples])
-        
+
         embedding_text = (
             f"Tool Name: {self.name}\n"
             f"Description: {self.description}\n"
@@ -106,10 +108,10 @@ def get_weather_forecast(location: str, days: int = 1):
     """Simulates fetching a weather forecast."""
     if not isinstance(location, str) or not isinstance(days, int):
         return {"error": "Invalid argument types. 'location' must be a string and 'days' an integer."}
-    
+
     weather_conditions = ["Sunny", "Cloudy", "Rainy", "Windy", "Snowy"]
     response = {"location": location, "forecast": []}
-    
+
     for i in range(days):
         date = (datetime.now() + timedelta(days=i)).strftime('%Y-%m-%d')
         condition = np.random.choice(weather_conditions)
@@ -272,23 +274,23 @@ def extract_arguments_hf(user_prompt: str, tool: Tool):
     chat = [
         {"role": "user", "content": f"{system_prompt}\n\nUser Prompt: \"{user_prompt}\""},
     ]
-    
+
     prompt = hf_tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
 
     try:
         inputs = hf_tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(hf_model.device)
-        
+
         # Generate with the model
         outputs = hf_model.generate(input_ids=inputs, max_new_tokens=256, do_sample=False)
         decoded_output = hf_tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
-        
+
         # Clean the response to find the JSON object
         json_str = decoded_output.strip()
-        
+
         # Find the first '{' and the last '}' to get the JSON part
         json_start = json_str.find('{')
         json_end = json_str.rfind('}')
-        
+
         if json_start != -1 and json_end != -1:
             json_str = json_str[json_start : json_end + 1]
             return json.loads(json_str)
@@ -302,7 +304,7 @@ def extract_arguments_hf(user_prompt: str, tool: Tool):
 def execute_tool(user_prompt: str):
     """The main pipeline: Find tool, extract args, execute."""
     selected_tool, score, _ = find_best_tool(user_prompt)
-    
+
     if USE_HF_LLM:
         print(f"⚙️ Selected Tool: {selected_tool.name}. Extracting arguments with Gemma...")
         extracted_args = extract_arguments_hf(user_prompt, selected_tool)
@@ -312,12 +314,17 @@ def execute_tool(user_prompt: str):
 
     if 'error' in extracted_args:
         print(f"❌ Argument extraction failed: {extracted_args['error']}")
+        # Ensure the final output string is valid JSON
+        final_output_str = json.dumps({
+            "error": "Execution failed during argument extraction.",
+            "details": extracted_args['error']
+        })
         return (
             user_prompt,
             selected_tool.name,
             f"{score:.3f}",
             json.dumps(extracted_args, indent=2),
-            "Execution failed during argument extraction."
+            final_output_str
         )
 
     print(f"✅ Arguments extracted: {json.dumps(extracted_args, indent=2)}")
@@ -349,7 +356,7 @@ def plot_tool_world(user_intent=None):
     tool_vectors = [tool.embedding.cpu().numpy() for tool in tools]
     labels = [tool.name for tool in tools]
     all_vectors = tool_vectors
-    
+
     if user_intent and user_intent.strip():
         intent_vector = embedder.encode(user_intent, convert_to_tensor=True).cpu().numpy()
         all_vectors.append(intent_vector)
@@ -361,7 +368,7 @@ def plot_tool_world(user_intent=None):
         n_neighbors = 1
 
     reducer = umap.UMAP(n_neighbors=n_neighbors, min_dist=0.3, metric='cosine', random_state=42)
-    
+
     # UMAP fit_transform requires at least 2 samples
     if len(all_vectors) < 2:
          # Create a dummy plot if there's not enough data
@@ -387,7 +394,7 @@ def plot_tool_world(user_intent=None):
     ax.set_xlabel("UMAP Dimension 1", fontsize=12)
     ax.set_ylabel("UMAP Dimension 2", fontsize=12)
     ax.grid(True)
-    
+
     handles, labels_legend = ax.get_legend_handles_labels()
     by_label = dict(zip(labels_legend, handles))
     ax.legend(by_label.values(), by_label.keys())
@@ -406,7 +413,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🛠️ Tool World: Advanced Prototype (Hugging Face Version)")
     gr.Markdown(
         "Enter a natural language command. The system will select the best tool, "
-        "extract structured arguments with **google/gemma-3b-it**, and execute it."
+        "extract structured arguments with **google/gemma-3n-E4B**, and execute it."
     )
 
     with gr.Row():
@@ -417,7 +424,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 lines=3
             )
             run_btn = gr.Button("Invoke Tool", variant="primary")
-            
+
             gr.Markdown("---")
             gr.Markdown("### Examples")
             gr.Examples(
@@ -435,7 +442,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             with gr.Row():
                 out_tool = gr.Textbox(label="Selected Tool", interactive=False)
                 out_score = gr.Textbox(label="Similarity Score", interactive=False)
-            
+
             out_args = gr.JSON(label="Extracted Arguments")
             out_result = gr.JSON(label="Tool Execution Output")
 
@@ -448,10 +455,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         if not user_prompt or not user_prompt.strip():
             # Return empty state and the default plot
             return "", "", {}, {}, plot_tool_world()
-        
+
         prompt, tool_name, score, args_json, result_json = execute_tool(user_prompt)
         fig = plot_tool_world(user_prompt)
-        
+
         # Safely load JSON strings into objects for the UI
         try:
             args_obj = json.loads(args_json)
@@ -470,7 +477,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         inputs=inp,
         outputs=[out_tool, out_score, out_args, out_result, plot_output]
     )
-    
+
     # Load the initial plot when the app starts
     demo.load(fn=lambda: plot_tool_world(None), inputs=None, outputs=plot_output)