import os, json, gradio as gr, torch
from transformers import AutoProcessor, AutoModelForCausalLM

hf_token = os.getenv("HF_TOKEN")
model_id = "google/functiongemma-270m-it"

# 1. Load for CPU specifically
processor = AutoProcessor.from_pretrained(model_id, token=hf_token)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.float32, # CPU prefers float32
    device_map={"": "cpu"},    # Forces everything onto CPU, avoiding "meta device"
    token=hf_token
)

def process_request(user_prompt, developer_prompt, tools_json):
    try:
        tools = json.loads(tools_json) if tools_json.strip() else []
        
        # FunctionGemma format
        messages = [
            {"role": "developer", "content": developer_prompt},
            {"role": "user", "content": user_prompt}
        ]

        # 2. Ensure inputs are on CPU
        inputs = processor.apply_chat_template(
            messages, tools=tools, add_generation_prompt=True, 
            return_dict=True, return_tensors="pt"
        ).to("cpu")

        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=128, do_sample=False)
        
        input_len = inputs.input_ids.shape[1]
        decoded = processor.decode(outputs[0][input_len:], skip_special_tokens=True)
        return decoded if decoded.strip() else "Model returned an empty string."
    
    except Exception as e:
        return f"Error: {str(e)}"

demo = gr.Interface(
    fn=process_request,
    inputs=[gr.Textbox(label="User Prompt"), gr.Textbox(label="Developer Prompt"), gr.Textbox(label="Tools (JSON Array)")],
    outputs=gr.Code(label="Model Output"),
    title="FunctionGemma CPU Fixed"
)

demo.launch(share=True)