Local model inference

README.md

@@ -15,7 +15,7 @@ executable APIs call given natural language instructions and API context.
 |gorilla-openfunctions-v0 | Given a function, and user intent, returns properly formatted json with the right arguments|
 |gorilla-openfunctions-v1 | + Parallel functions, and can choose between functions|
 
-## Example Usage
+## Example Usage (Hosted)
 
 1. OpenFunctions is compatible with OpenAI Functions
 
@@ -63,6 +63,75 @@ get_gorilla_response(query, functions=functions)
 ```bash
 uber.ride(loc="berkeley", type="plus", time=10)
 ```
+
+## Example Usage (Run Locally)
+
+```python
+import json
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+def get_prompt(user_query: str, functions: list = []) -> str:
+    """
+    Generates a conversation prompt based on the user's query and a list of functions.
+
+    Parameters:
+    - user_query (str): The user's query.
+    - functions (list): A list of functions to include in the prompt.
+
+    Returns:
+    - str: The formatted conversation prompt.
+    """
+    if len(functions) == 0:
+        return f"USER: <<question>> {user_query}\nASSISTANT: "
+    functions_string = json.dumps(functions)
+    return f"USER: <<question>> {user_query} <<function>> {functions_string}\nASSISTANT: "
+
+# Device setup
+device : str = "cuda:0" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+# Model and tokenizer setup
+model_id : str = "gorilla-llm/gorilla-openfunctions-v0"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True)
+
+# Move model to device
+model.to(device)
+
+# Pipeline setup
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=128,
+    batch_size=16,
+    torch_dtype=torch_dtype,
+    device=device,
+)
+
+# Example usage
+query: str = "Call me an Uber ride type \"Plus\" in Berkeley at zipcode 94704 in 10 minutes"
+functions = [
+    {
+        "name": "Uber Carpool",
+        "api_name": "uber.ride",
+        "description": "Find suitable ride for customers given the location, type of ride, and the amount of time the customer is willing to wait as parameters",
+        "parameters":  [
+            {"name": "loc", "description": "Location of the starting place of the Uber ride"},
+            {"name": "type", "enum": ["plus", "comfort", "black"], "description": "Types of Uber ride user is ordering"},
+            {"name": "time", "description": "The amount of time in minutes the customer is willing to wait"}
+        ]
+    }
+]
+
+# Generate prompt and obtain model output
+prompt = get_prompt(query, functions=functions)
+output = pipe(prompt)
+
+print(output)
+```
+
    
 ## Contributing