Update app.py

app.py

@@ -0,0 +1,151 @@
+# app.py
+import json
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Load model and tokenizer
+model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+# Define a simple addition function schema
+function_schema = {
+    "name": "add_numbers",
+    "description": "Add two numbers together",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "number1": {
+                "type": "number",
+                "description": "The first number"
+            },
+            "number2": {
+                "type": "number",
+                "description": "The second number"
+            }
+        },
+        "required": ["number1", "number2"]
+    }
+}
+
+# Create prompt with function definition
+def create_prompt(user_input, function):
+    prompt = f"<|system|>\nYou are a helpful assistant that can use functions. Please call the add_numbers function for any addition requests.\n\nAvailable function:\n{json.dumps(function)}\n<|user|>\n{user_input}\n<|assistant|>\n"
+    return prompt
+
+# Extract function call from response
+def extract_function_call(response_text):
+    try:
+        if "<functioncall>" in response_text and "</functioncall>" in response_text:
+            func_text = response_text.split("<functioncall>")[1].split("</functioncall>")[0].strip()
+            return json.loads(func_text)
+        return None
+    except Exception as e:
+        print(f"Error extracting function call: {e}")
+        return None
+
+# Actually perform the addition
+def execute_add_numbers(params):
+    try:
+        num1 = float(params.get("number1", 0))
+        num2 = float(params.get("number2", 0))
+        return {"result": num1 + num2}
+    except Exception as e:
+        return {"error": str(e)}
+
+def process_query(query, debug=False):
+    # Create the initial prompt
+    prompt = create_prompt(query, function_schema)
+    
+    # Generate the initial response
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        temperature=0.1
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    
+    # Process the response
+    try:
+        assistant_response = response.split("<|assistant|>")[1].strip()
+    except:
+        return "Error parsing model response."
+    
+    debug_info = f"Initial response:\n{assistant_response}\n\n" if debug else ""
+    
+    # Check for function call
+    function_call = extract_function_call(assistant_response)
+    if not function_call:
+        return debug_info + "No function call detected in the response."
+    
+    debug_info += f"Function call detected:\n{json.dumps(function_call, indent=2)}\n\n" if debug else ""
+    
+    # Execute the function
+    result = execute_add_numbers(function_call)
+    
+    debug_info += f"Function result:\n{json.dumps(result, indent=2)}\n\n" if debug else ""
+    
+    # Create follow-up prompt with function result
+    follow_up_prompt = f"{prompt}\n<functioncall>\n{json.dumps(function_call)}\n</functioncall>\n\n<functionresponse>\n{json.dumps(result)}\n</functionresponse>\n"
+    
+    # Generate final response
+    follow_up_inputs = tokenizer(follow_up_prompt, return_tensors="pt").to(model.device)
+    follow_up_outputs = model.generate(
+        **follow_up_inputs,
+        max_new_tokens=256,
+        temperature=0.1
+    )
+    follow_up_response = tokenizer.decode(follow_up_outputs[0], skip_special_tokens=False)
+    
+    try:
+        if "<functionresponse>" in follow_up_response and "</functionresponse>" in follow_up_response:
+            final_response = follow_up_response.split("</functionresponse>")[1].strip()
+        else:
+            final_response = follow_up_response.split("<|assistant|>")[1].strip()
+    except:
+        return debug_info + "Error extracting final response."
+    
+    if debug:
+        return debug_info + f"Final response:\n{final_response}"
+    else:
+        return final_response
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Llama 3.1 Function Calling: Addition Calculator")
+    gr.Markdown("Ask the model to add numbers, and it will use the `add_numbers` function")
+    
+    with gr.Row():
+        query_input = gr.Textbox(
+            label="Your Question",
+            placeholder="Example: What is 24 plus 18?",
+            lines=2
+        )
+        debug_checkbox = gr.Checkbox(label="Show Debug Info", value=False)
+    
+    submit_btn = gr.Button("Submit")
+    
+    output = gr.Textbox(label="Response", lines=10)
+    
+    submit_btn.click(
+        fn=process_query,
+        inputs=[query_input, debug_checkbox],
+        outputs=output
+    )
+    
+    gr.Examples(
+        [
+            ["What is 25 plus 17?"],
+            ["Can you add 123 and 456?"],
+            ["Calculate 3.14 + 2.71"]
+        ],
+        inputs=query_input
+    )
+
+demo.launch()