Fix: fix FLAN-XXL input/output format

Author: Shing Lyu

kendra_retriever_samples/kendra_chat_flan_xxl.py

@@ -30,12 +30,13 @@ class ContentHandler(LLMContentHandler):
       accepts = "application/json"
 
       def transform_input(self, prompt: str, model_kwargs: dict) -> bytes:
-          input_str = json.dumps({"inputs": prompt, "parameters": model_kwargs})
+          input_str = json.dumps({"text_inputs": prompt, **model_kwargs})
           return input_str.encode('utf-8')
 
       def transform_output(self, output: bytes) -> str:
           response_json = json.loads(output.read().decode("utf-8"))
-          return response_json[0]["generated_text"]
+          print(response_json)
+          return response_json["generated_texts"][0]
 
   content_handler = ContentHandler()
 

kendra_retriever_samples/kendra_retriever_flan_xxl.py

@@ -18,12 +18,12 @@ class ContentHandler(LLMContentHandler):
         accepts = "application/json"
 
         def transform_input(self, prompt: str, model_kwargs: dict) -> bytes:
-            input_str = json.dumps({"inputs": prompt, "parameters": model_kwargs})
+            input_str = json.dumps({"text_inputs": prompt, **model_kwargs})
             return input_str.encode('utf-8')
 
         def transform_output(self, output: bytes) -> str:
             response_json = json.loads(output.read().decode("utf-8"))
-            return response_json[0]["generated_text"]
+            return response_json["generated_texts"][0]
 
     content_handler = ContentHandler()