t5 + flant5

models.py

@@ -8,12 +8,12 @@ auth_token = os.environ.get("CLARIN_KNEXT")
 
 DEFAULT_INPUTS: Dict[str, str] = {
     "polish": (
-        "[U] Chciałbym zarezerwować stolik na 4 osoby na piątek o godzinie 18.30. "
+        "[U] Chciałbym zarezerwować stolik na 4 osoby na piątek o godzinie 18:30. "
         "[Dziedzina] Restauracje: Popularna usługa wyszukiwania i rezerwacji restauracji "
         "[Atrybut] Czas: Wstępny czas rezerwacji restauracji"
     ),
     "english": (
-        "[U] I want to book a table for 4 people on Friday, 6.30 pm. "
+        "[U] I want to book a table for 4 people on Friday, 6:30 pm. "
         "[Domain] Restaurants: A popular restaurant search and reservation service "
         "[Slot] Time: Tentative time of restaurant reservation"
     ),
@@ -30,6 +30,26 @@ MODELS: Dict[str, Dict[str, Any]] = {
         "tokenizer": T5Tokenizer.from_pretrained("clarin-knext/plt5-base-dst", use_auth_token=auth_token),
         "default_input": DEFAULT_INPUTS["polish"],
     },
+    "t5-small": {
+        "model": T5ForConditionalGeneration.from_pretrained("clarin-knext/t5-small-dst", use_auth_token=auth_token),
+        "tokenizer": T5Tokenizer.from_pretrained("clarin-knext/t5-small-dst", use_auth_token=auth_token),
+        "default_input": DEFAULT_INPUTS["english"],
+    },
+    "t5-base": {
+        "model": T5ForConditionalGeneration.from_pretrained("clarin-knext/t5-base-dst", use_auth_token=auth_token),
+        "tokenizer": T5Tokenizer.from_pretrained("clarin-knext/t5-base-dst", use_auth_token=auth_token),
+        "default_input": DEFAULT_INPUTS["english"],
+    },
+    "flant5-small": {
+        "model": T5ForConditionalGeneration.from_pretrained("clarin-knext/flant5-small-dst", use_auth_token=auth_token),
+        "tokenizer": T5Tokenizer.from_pretrained("clarin-knext/flant5-small-dst", use_auth_token=auth_token),
+        "default_input": DEFAULT_INPUTS["polish"],
+    },
+    "flant5-base": {
+        "model": T5ForConditionalGeneration.from_pretrained("clarin-knext/flant5-base-dst", use_auth_token=auth_token),
+        "tokenizer": T5Tokenizer.from_pretrained("clarin-knext/flant5-base-dst", use_auth_token=auth_token),
+        "default_input": DEFAULT_INPUTS["polish"],
+    },
 }
 
 PIPELINES: Dict[str, Pipeline] = {