Allow Any inputs

pkg/cli/predict.go

@@ -279,7 +279,7 @@ func handleMultipleFileOutput(prediction *predict.Response, outputSchema *openap
 }
 
 func parseInputFlags(inputs []string) (predict.Inputs, error) {
-	keyVals := map[string]string{}
+	keyVals := map[string][]string{}
 	for _, input := range inputs {
 		var name, value string
 
@@ -295,8 +295,12 @@ func parseInputFlags(inputs []string) (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

@@ -13,50 +13,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
@@ -67,6 +59,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

@@ -24,7 +24,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

@@ -40,9 +40,6 @@
     Path as CogPath,
 )
 
-ALLOWED_INPUT_TYPES = [str, int, float, bool, CogFile, CogPath]
-
-
 class BasePredictor(ABC):
     def setup(self, weights: Optional[Union[CogFile, CogPath]] = None) -> None:
         """
@@ -185,6 +182,8 @@ class Config:
         # But, after validation, we want to pass the actual value to predict(), not the enum object
         use_enum_values = True
 
+        arbitrary_types_allowed = True
+
     def cleanup(self) -> None:
         """
         Cleanup any temporary files created by the input.
@@ -231,15 +230,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

@@ -3,6 +3,7 @@
 
 import pytest
 import responses
+from fastapi.exceptions import FastAPIError
 
 from .conftest import make_client, uses_predictor
 
@@ -213,10 +214,5 @@ def test_choices_int(client):
 
 
 def test_untyped_inputs():
-    with pytest.raises(TypeError):
+    with pytest.raises(FastAPIError):
         make_client("input_untyped")
-
-
-def test_input_with_unsupported_type():
-    with pytest.raises(TypeError):
-        make_client("input_unsupported_type")