Update app.py

app.py

@@ -4,43 +4,52 @@ from transformers import AutoProcessor, AutoModelForCausalLM
 hf_token = os.getenv("HF_TOKEN")
 model_id = "google/functiongemma-270m-it"
 
+# Load model
 processor = AutoProcessor.from_pretrained(model_id, token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
-    model_id, torch_dtype=torch.float16, device_map="auto", token=hf_token
+    model_id, 
+    torch_dtype=torch.float16, 
+    device_map="auto", 
+    token=hf_token,
+    low_cpu_mem_usage=True
 )
 
 def process_request(user_prompt, tools_json):
     try:
-        # 1. Parse the tools sent from the client
         tools = json.loads(tools_json) if tools_json.strip() else []
-        
-        # 2. Build the message history
         messages = [
             {"role": "developer", "content": "You are a model that can do function calling with the following functions"},
             {"role": "user", "content": user_prompt}
         ]
-
-        # 3. Apply template with DYNAMIC tools
         inputs = processor.apply_chat_template(
             messages, tools=tools, add_generation_prompt=True, 
             return_dict=True, return_tensors="pt"
         ).to(model.device)
 
-        # 4. Generate
         with torch.no_grad():
             outputs = model.generate(**inputs, max_new_tokens=128, do_sample=False)
         
-        # 5. Decode just the new part
         input_len = inputs.input_ids.shape[1]
         return processor.decode(outputs[0][input_len:], skip_special_tokens=True)
-    
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Gradio interface with TWO inputs
-demo = gr.Interface(
-    fn=process_request,
-    inputs=[gr.Textbox(label="User Prompt"), gr.Textbox(label="Tools (JSON Array)")],
-    outputs=gr.Code(label="Model Output"),
-)
-demo.launch()
+# Create the UI using Blocks for better API support
+with gr.Blocks() as demo:
+    gr.Markdown("# FunctionGemma API Server")
+    with gr.Row():
+        prompt_input = gr.Textbox(label="User Prompt")
+        tools_input = gr.Textbox(label="Tools (JSON Array)")
+    output_text = gr.Code(label="Model Output")
+    
+    submit_btn = gr.Button("Submit")
+    
+    # CRITICAL: This 'api_name' must match what your client expects
+    submit_btn.click(
+        fn=process_request, 
+        inputs=[prompt_input, tools_input], 
+        outputs=output_text,
+        api_name="predict" 
+    )
+
+demo.launch()