allow arbitrary input types for model inputs

pkg/cli/predict.go

@@ -258,7 +258,7 @@ func handleMultipleFileOutput(prediction *predict.Response, outputSchema *openap
 
 func parseInputFlags(inputs []string, schema *openapi3.T) (predict.Inputs, error) {
 	var err error
-	keyVals := map[string]string{}
+	keyVals := map[string][]string{}
 	for _, input := range inputs {
 		var name, value string
 
@@ -277,7 +277,12 @@ func parseInputFlags(inputs []string, schema *openapi3.T) (predict.Inputs, error
 		if strings.HasPrefix(value, `"`) && strings.HasSuffix(value, `"`) {
 			value = value[1 : len(value)-1]
 		}
-		keyVals[name] = value
+		currVal, present := keyVals[name]
+		if present {
+			keyVals[name] = append(currVal, value)
+		} else {
+			keyVals[name] = []string{value}
+		}
 	}
 	return predict.NewInputs(keyVals), nil
 }

pkg/predict/input.go

@@ -12,50 +12,42 @@ import (
 )
 
 type Input struct {
-	String *string
-	File   *string
+	String     *string
+	File       *string
+	StringList *[]string
 }
 
 type Inputs map[string]Input
 
-func NewInputs(keyVals map[string]string) Inputs {
+func NewInputs(keyVals map[string][]string) Inputs {
 	input := Inputs{}
-	for key, val := range keyVals {
-		val := val
-		if strings.HasPrefix(val, "@") {
-			val = val[1:]
-			expandedVal, err := homedir.Expand(val)
-			if err != nil {
-				// FIXME: handle this better?
-				console.Warnf("Error expanding homedir: %s", err)
+	for key, arr := range keyVals {
+		arr := arr
+		// Handle singleton slices by converting them to strings or expanding filenames
+		if len(arr) == 1 {
+			val := arr[0]
+			if strings.HasPrefix(val, "@") {
+				val = val[1:]
+				expandedVal, err := homedir.Expand(val)
+				if err != nil {
+					// FIXME: handle this better?
+					console.Warnf("Error expanding homedir: %s", err)
+				} else {
+					val = expandedVal
+				}
+				input[key] = Input{File: &val}
 			} else {
-				val = expandedVal
+				input[key] = Input{String: &val}
 			}
-
-			input[key] = Input{File: &val}
-		} else {
-			input[key] = Input{String: &val}
-		}
-	}
-	return input
-}
-
-func NewInputsWithBaseDir(keyVals map[string]string, baseDir string) Inputs {
-	input := Inputs{}
-	for key, val := range keyVals {
-		val := val
-		if strings.HasPrefix(val, "@") {
-			val = filepath.Join(baseDir, val[1:])
-			input[key] = Input{File: &val}
 		} else {
-			input[key] = Input{String: &val}
+			input[key] = Input{StringList: &arr}
 		}
 	}
 	return input
 }
 
-func (inputs *Inputs) toMap() (map[string]string, error) {
-	keyVals := map[string]string{}
+func (inputs *Inputs) toMap() (map[string]interface{}, error) {
+	keyVals := map[string]interface{}{}
 	for key, input := range *inputs {
 		if input.String != nil {
 			keyVals[key] = *input.String
@@ -66,6 +58,8 @@ func (inputs *Inputs) toMap() (map[string]string, error) {
 			}
 			mimeType := mime.TypeByExtension(filepath.Ext(*input.File))
 			keyVals[key] = dataurl.New(content, mimeType).String()
+		} else if input.StringList != nil {
+			keyVals[key] = *input.StringList
 		}
 	}
 	return keyVals, nil

pkg/predict/predictor.go

@@ -23,7 +23,7 @@ type HealthcheckResponse struct {
 
 type Request struct {
 	// TODO: could this be Inputs?
-	Input map[string]string `json:"input"`
+	Input map[string]interface{} `json:"input"`
 }
 
 type Response struct {

python/cog/predictor.py

@@ -25,9 +25,6 @@
 )
 
 
-ALLOWED_INPUT_TYPES = [str, int, float, bool, CogFile, CogPath]
-
-
 class BasePredictor(ABC):
     def setup(self, weights: Optional[Union[CogFile, CogPath]] = None) -> None:
         """
@@ -213,15 +210,6 @@ class Input(BaseModel):
     for name, parameter in signature.parameters.items():
         InputType = parameter.annotation
 
-        if InputType is inspect.Signature.empty:
-            raise TypeError(
-                f"No input type provided for parameter `{name}`. Supported input types are: {readable_types_list(ALLOWED_INPUT_TYPES)}."
-            )
-        elif InputType not in ALLOWED_INPUT_TYPES:
-            raise TypeError(
-                f"Unsupported input type {human_readable_type_name(InputType)} for parameter `{name}`. Supported input types are: {readable_types_list(ALLOWED_INPUT_TYPES)}."
-            )
-
         # if no default is specified, create an empty, required input
         if parameter.default is inspect.Signature.empty:
             default = Input()

python/tests/server/test_http_input.py

@@ -103,7 +103,7 @@ def test_file_input_with_http_url(client, httpserver, match):
 
 @uses_predictor("input_path_2")
 def test_file_input_with_http_url_error(client, httpserver, match):
-    httpserver.expect_request("/foo.txt").respond_with_data("haha", status = 404)
+    httpserver.expect_request("/foo.txt").respond_with_data("haha", status=404)
     resp = client.post(
         "/predictions",
         json={"input": {"path": httpserver.url_for("/foo.txt")}},
@@ -216,8 +216,3 @@ def test_choices_int(client):
 def test_untyped_inputs():
     with pytest.raises(TypeError):
         make_client("input_untyped")
-
-
-def test_input_with_unsupported_type():
-    with pytest.raises(TypeError):
-        make_client("input_unsupported_type")