Fix Client and serving models

MonsterAPIClient.py

@@ -96,9 +96,9 @@ class MClient():
             'Authorization': 'Bearer ' + self.auth_token}
         self.base_url = 'https://api.monsterapi.ai/v1'
         self.models_to_data_model = MODELS_TO_DATAMODEL
-        self.mock = os.environ.get('MOCK_Runner', "True").lower() == "true"
+        self.mock = os.environ.get('MOCK_Runner', "False").lower() == "true"
 
-    def get_response(self, model:Literal['falcon-20b-instruct', 'falcon-7b-instruct', 'mpt-30B-instruct', 'mpt-7B-instruct'], 
+    def get_response(self, model:Literal['falcon-7b-instruct', 'falcon-40b-instruct', 'mpt-30B-instruct', 'mpt-7b-instruct', 'openllama-13b-base', 'llama2-7b-chat'], 
                      data: dict):
     
         if model not in self.models_to_data_model:
@@ -106,8 +106,9 @@ class MClient():
 
         dataModel = self.models_to_data_model[model](**data)
         url = f"{self.base_url}/generate/{model}"
-        #url = self.base_url + model
         data = dataModel.dict()
+        logger.info(f"Calling Monster API with url: {url}, with payload: {data}")
+
         # convert all values into string
         for key, value in data.items():
             data[key] = str(value)
@@ -123,7 +124,7 @@ class MClient():
         response.raise_for_status()
         return response.json()
     
-    def wait_and_get_result(self, process_id, timeout=300):
+    def wait_and_get_result(self, process_id, timeout=100):
         start_time = time.time()
         while True:
             elapsed_time = time.time() - start_time
@@ -139,12 +140,12 @@ class MClient():
             else:
                 if self.mock:
                     return 100 * "Mock Output!"
-                logger.info(f"Process {process_id} is still running, status is {status['status']}. Waiting for 5 seconds...")
-                time.sleep(0.05)
+                logger.info(f"Process {process_id} is still running, status is {status['status']}. Waiting ...")
+                time.sleep(0.01)
 
         
 if __name__ == '__main__':
     client = MClient()
-    response = client.get_response('falcon-7b-instruct', {"prompt": 'How to make a sandwich'})
+    response = client.get_response('falcon-7b-instruct', {"prompt": 'How to make a sandwich?'})
     output = client.wait_and_get_result(response['process_id'])
     print(output)

gradio_app.py

@@ -7,8 +7,7 @@ from MonsterAPIClient import MODELS_TO_DATAMODEL
 client = MClient()
 
 # Available models list
-EXCLUSION_LIST = ['mpt-30B-instruct', 'llama2-7b-chat', 'openllama-13b-base']
-available_models = list(set(list(MODELS_TO_DATAMODEL.keys())) - set(EXCLUSION_LIST))
+MODELS_TO_SERVE =  ['llama2-7b-chat','mpt-7b-instruct','falcon-7b-instruct']
 
 def generate_model_output(model: str, input_text: str, temp: float = 0.98) -> str:
     """
@@ -28,12 +27,12 @@ def generate_model_output(model: str, input_text: str, temp: float = 0.98) -> st
             "temp": temp,
         })
         output = client.wait_and_get_result(response['process_id'])
-        return model, output
+        return model, output['text']
     except Exception as e:
         return model, f"Error occurred: {str(e)}"
 
 def generate_output(selected_models: list, input_text: str, temp: float = 0.98,
-                    available_models: list = available_models) -> list:
+                    available_models: list = MODELS_TO_SERVE) -> list:
     """
     Generate outputs from selected models using Monster API.
 
@@ -58,13 +57,13 @@ def generate_output(selected_models: list, input_text: str, temp: float = 0.98,
         if model not in outputs:
             ret_outputs.append("Model not selected!")
         else:
-            ret_outputs.append(outputs[model].replace("\n", "<br>"))
+            ret_outputs.append(outputs[model])
 
-    return ret_outputs
+    return ret_outputs  
 
-output_components = [gr.outputs.Textbox(label=model) for model in available_models]
+output_components = [gr.outputs.Textbox(label=model) for model in MODELS_TO_SERVE]
 
-checkboxes = gr.inputs.CheckboxGroup(available_models, label="Select models to generate outputs:")
+checkboxes = gr.inputs.CheckboxGroup(MODELS_TO_SERVE, label="Select models to generate outputs:")
 textbox = gr.inputs.Textbox(label="Input Prompt")
 temp = gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.98, label="Temperature", step=0.01)