import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载 xLAM 模型和 tokenizer
model_name = "Salesforce/xLAM-7b-r"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 定义任务提示和格式提示
task_instruction = """
Based on the previous context and API request history, generate an API request or a response as an AI assistant.
""".strip()

format_instruction = """
The output should be of the JSON format, which specifies a list of generated function calls. If no function call is needed, please make tool_calls an empty list "[]".
""".strip()

# 定义工具信息
get_weather_api = {
    "name": "get_weather",
    "description": "Get the current weather for a location",
    "parameters": {
        "type": "object",
        "properties": {
            "location": {
                "type": "string",
                "description": "The city and state, e.g. San Francisco, New York"
            },
            "unit": {
                "type": "string",
                "enum": ["celsius", "fahrenheit"],
                "description": "The unit of temperature to return"
            }
        },
        "required": ["location"]
    }
}

search_api = {
    "name": "search",
    "description": "Search for information on the internet",
    "parameters": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "The search query, e.g. 'latest news on AI'"
            }
        },
        "required": ["query"]
    }
}

# 转换工具为 xLAM 的格式
def convert_to_xlam_tool(tools):
    if isinstance(tools, dict):
        return {
            "name": tools["name"],
            "description": tools["description"],
            "parameters": {k: v for k, v in tools["parameters"].get("properties", {}).items()}
        }
    elif isinstance(tools, list):
        return [convert_to_xlam_tool(tool) for tool in tools]
    else:
        return tools

xlam_format_tools = convert_to_xlam_tool([get_weather_api, search_api])

# 生成提示
def build_prompt(task_instruction, format_instruction, tools, query):
    prompt = f"[BEGIN OF TASK INSTRUCTION]\n{task_instruction}\n[END OF TASK INSTRUCTION]\n\n"
    prompt += f"[BEGIN OF AVAILABLE TOOLS]\n{tools}\n[END OF AVAILABLE TOOLS]\n\n"
    prompt += f"[BEGIN OF FORMAT INSTRUCTION]\n{format_instruction}\n[END OF FORMAT INSTRUCTION]\n\n"
    prompt += f"[BEGIN OF QUERY]\n{query}\n[END OF QUERY]\n\n"
    return prompt

# 定义模型推理函数
def generate_response(query):
    # 构建输入提示
    content = build_prompt(task_instruction, format_instruction, xlam_format_tools, query)
    messages = [{'role': 'user', 'content': content}]
    
    # 编码输入
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
    
    # 生成输出
    outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
    
    # 解码输出
    response = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
    return response

# 使用 Gradio 创建简单的 Web 应用
with gr.Blocks() as demo:
    gr.Markdown("## 使用 xLAM 模型进行智能对话")
    
    query = gr.Textbox(label="输入您的问题", placeholder="请输入您的问题")
    output = gr.Textbox(label="模型响应")
    
    submit_btn = gr.Button("提交")
    submit_btn.click(fn=generate_response, inputs=query, outputs=output)

# 启动 Gradio 应用
demo.launch()