update

main.py

@@ -5,38 +5,35 @@ import json
 
 from transformers import pipeline
 
-app = FastAPI()
+def get_prompt(user_query: str, functions: list = []) -> str:
+    """
+    Generates a conversation prompt based on the user's query and a list of functions.
 
-@app.post("/query_gorilla")
-async def query_gorilla(req: Request):
-    body = await req.body()
-    parsedBody = json.loads(body)
-    print(parsedBody['query'])
-    print(parsedBody['functions'])
-    return {
-      "val": body
-    }
+    Parameters:
+    - user_query (str): The user's query.
+    - functions (list): A list of functions to include in the prompt.
 
-app.mount("/", StaticFiles(directory="static", html=True), name="static")
-
-@app.get("/")
-def index() -> FileResponse:
-    return FileResponse(path="/app/static/index.html", media_type="text/html")
+    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: "
 
-TODO = '''
-print('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
 
-print('Model and tokenizer setup')
 model_id : str = "gorilla-llm/gorilla-openfunctions-v1"
+print('AutoTokenizer.from_pretrained ...')
 tokenizer = AutoTokenizer.from_pretrained(model_id)
+print('AutoModelForCausalLM.from_pretrained ...')
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True)
 
-print('Move model to device')
+print('mode.to(device) ...')
 model.to(device)
 
-print('Pipeline setup')
+print('Pipeline setup ...')
 pipe = pipeline(
     "text-generation",
     model=model,
@@ -46,5 +43,29 @@ pipe = pipeline(
     torch_dtype=torch_dtype,
     device=device,
 )
-'''
+
+print('FastAPI setup ...')
+app = FastAPI()
+
+@app.post("/query_gorilla")
+async def query_gorilla(req: Request):
+    body = await req.body()
+    parsedBody = json.loads(body)
+    print(parsedBody['query'])
+    print(parsedBody['functions'])
+
+    print('Generate prompt and obtain model output')
+    prompt = get_prompt(query, functions=functions)
+    output = pipe(prompt)
+    
+    return {
+      "val": output
+    }
+
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
+
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(path="/app/static/index.html", media_type="text/html")
+