Merge pull request #394 from UKP-SQuARE/skills

Pass all kwargs to models

Author: timbmg

evaluator/requirements.txt

@@ -1,5 +1,5 @@
 square-elk-json-formatter==0.0.3
-square-skill-api==0.0.35
+square-skill-api==0.0.37
 square-auth==0.0.14
 uvicorn>=0.15.0
 fastapi>=0.70.0

skill-manager/requirements.txt

@@ -1,5 +1,5 @@
 square-elk-json-formatter==0.0.3
-square-skill-api==0.0.35
+square-skill-api==0.0.37
 square-auth==0.0.14
 uvicorn>=0.15.0
 fastapi>=0.70.0

skill-manager/skill_manager/models.py

@@ -65,7 +65,7 @@ class Skill(MongoModel):
         description="A description of the skill, for example describing its pipeline.",
     )
     default_skill_args: Optional[Dict] = Field(
-        None,
+        {},
         description="A dictionary holding key-value pairs that should always be sent to the skill as input. This allows to use the same skill implementataion in different ways.",
     )
     published: bool = Field(

skill-manager/skill_manager/routers/skill.py

@@ -6,7 +6,6 @@
 from threading import Thread
 from typing import Dict, List
 
-import requests_cache
 from bson import ObjectId
 from fastapi import APIRouter, Depends, Request
 from square_auth.auth import Auth
@@ -25,6 +24,7 @@
 from skill_manager.keycloak_api import KeycloakAPI
 from skill_manager.models import Prediction, Skill, SkillType
 from skill_manager.routers import client_credentials
+from skill_manager.utils import merge_dicts
 
 logger = logging.getLogger(__name__)
 
@@ -236,13 +236,29 @@ async def query_skill(
     query = query_request.query
     user_id = query_request.user_id
 
-    skill = await get_skill_if_authorized(request, skill_id=id, write_access=False)
+    skill: Skill = await get_skill_if_authorized(
+        request, skill_id=id, write_access=False
+    )
     query_request.skill = json.loads(skill.json())
 
-    default_skill_args = skill.default_skill_args
-    if default_skill_args is not None:
-        # add default skill args, potentially overwrite with query.skill_args
-        query_request.skill_args = {**default_skill_args, **query_request.skill_args}
+    # merge kargs with kwargs in request
+    for kwargs_key in [
+        "model_kwargs",
+        "task_kwargs",
+        "preprocessing_kwargs",
+        "explain_kwargs",
+        "attack_kwargs",
+    ]:
+        # overwrite kwargs from query_request with default kwargs
+        kwargs = merge_dicts(
+            skill.default_skill_args.pop(kwargs_key, {}),
+            getattr(query_request, kwargs_key),
+        )
+        # set kwargs in query_request
+        setattr(query_request, kwargs_key, kwargs)
+    query_request.skill_args = merge_dicts(
+        skill.default_skill_args, query_request.skill_args
+    )
 
     headers = {"Authorization": f"Bearer {token}"}
     if request.headers.get("Cache-Control"):

skill-manager/skill_manager/utils.py

@@ -0,0 +1,6 @@
+def merge_dicts(*dicts):
+    """Merge multiple dictionaries into one. Overwrites values from left to right."""
+    merged = {}
+    for d in dicts:
+        merged.update(d)
+    return merged

skill-manager/tests/conftest.py

@@ -92,7 +92,7 @@ def skill_init(
             user_id=user_id,
             description=description,
             published=published,
-            default_skill_args=default_skill_args,
+            default_skill_args={} if default_skill_args is None else default_skill_args,
             **kwargs,
         )
         if not skill.id:

skill-manager/tests/test_api.py

@@ -503,7 +503,11 @@ def test_query_skill_with_default_skill_args(
 ):
     test_realm = "test-realm"
     test_user = "test-user"
-    default_skill_args = {"adapter": "my-adapter", "context": "default context"}
+    default_skill_args = {
+        "adapter": "my-adapter",
+        "context": "default context",
+        "model_kwargs": {"model_foo": "model_bar"},
+    }
     response, skill = create_skill_via_api(
         pers_client,
         token_factory,
@@ -540,10 +544,19 @@ def test_query_skill_with_default_skill_args(
         headers=dict(Authorization="Bearer " + token),
     )
 
-    actual_skill_query_body = json.loads(responses.calls[0].request.body)["skill_args"]
+    actual_request_body = json.loads(responses.calls[0].request.body)
+
+    # model_kwargs is supposed to be removed from the skill_args and parsed separately
+    TestCase().assertDictEqual(
+        actual_request_body["model_kwargs"], default_skill_args.pop("model_kwargs")
+    )
+
+    # remaining args should end up in skill_args
     expected_skill_query_body = default_skill_args
     expected_skill_query_body["context"] = query_context["context"]
-    TestCase().assertDictEqual(actual_skill_query_body, expected_skill_query_body)
+    TestCase().assertDictEqual(
+        actual_request_body["skill_args"], expected_skill_query_body
+    )
 
 
 @responses.activate

skill-manager/tests/test_utils.py

@@ -0,0 +1,15 @@
+from skill_manager.utils import merge_dicts
+
+
+def test_merge_dicts():
+    d1 = {"a": 1}
+    d2 = {"b": 2}
+    merged_dicts = merge_dicts(d1, d2)
+    assert merged_dicts == {"a": 1, "b": 2}
+
+
+def test_overwrite_merge_dicts():
+    d1 = {"a": 1}
+    d2 = {"a": 2}
+    merged_dicts = merge_dicts(d1, d2)
+    assert merged_dicts == {"a": 2}

skills/Dockerfile

@@ -12,6 +12,7 @@ COPY requirements.txt ./
 RUN pip install -r requirements.txt
 
 COPY main.py main.py
+COPY utils.py utils.py
 ARG skill
 COPY ./$skill/skill.py skill.py 
 COPY logging.conf logging.conf

skills/boolq/skill.py

@@ -3,6 +3,8 @@
 from square_model_client import SQuAREModelClient
 from square_skill_api.models import QueryOutput, QueryRequest
 
+from utils import extract_model_kwargs_from_request
+
 logger = logging.getLogger(__name__)
 
 square_model_client = SQuAREModelClient()
@@ -12,18 +14,15 @@ async def predict(request: QueryRequest) -> QueryOutput:
     """Predicts yes/no for a boolean question with context"""
     query = request.query
     context = request.skill_args["context"]
-    explain_kwargs = request.explain_kwargs or {}
-    attack_kwargs = request.attack_kwargs or {}
+
+    model_request_kwargs = extract_model_kwargs_from_request(request)
 
     prepared_input = [[context, query]]
 
     model_request = {
         "input": prepared_input,
-        "preprocessing_kwargs": {},
-        "model_kwargs": {},
         "adapter_name": request.skill_args["adapter"],
-        "explain_kwargs": explain_kwargs,
-        "attack_kwargs": attack_kwargs,
+        **model_request_kwargs,
     }
     model_response = await square_model_client(
         model_name=request.skill_args["base_model"],