bug fix

app.py

@@ -1,25 +1,46 @@
-import gradio as gr
-import torch
+import os, json, gradio as gr, torch
 from transformers import AutoProcessor, AutoModelForCausalLM
 
-# Load the model and processor
+hf_token = os.getenv("HF_TOKEN")
 model_id = "google/functiongemma-270m-it"
-processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
 
-def generate_tool_call(user_prompt):
-    # This model is designed for function calling/tool use
-    # It requires a specific prompt format (see model card for details)
-    inputs = processor(text=user_prompt, return_tensors="pt").to(model.device)
-    output = model.generate(**inputs, max_new_tokens=128)
-    return processor.decode(output[0], skip_special_tokens=True)
+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
+)
+
+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=generate_tool_call,
-    inputs="text",
-    outputs="text",
-    title="FunctionGemma 270M Demo",
-    description="A lightweight 270M model specialized in function calling."
+    fn=process_request,
+    inputs=[gr.Textbox(label="User Prompt"), gr.Textbox(label="Tools (JSON Array)")],
+    outputs=gr.Code(label="Model Output"),
 )
-
 demo.launch()