Update server.py

server.py

@@ -48,14 +48,43 @@ class SplitArgs(argparse.Action):
         setattr(
             namespace, self.dest, values.replace('"', "").replace("'", "").split(",")
         )
+
+
+# Script arguments
 parser = argparse.ArgumentParser(
     prog="TavernAI Extras", description="Web API for transformers models"
 )
+parser.add_argument(
+    "--port", type=int, help="Specify the port on which the application is hosted"
+)
+parser.add_argument(
+    "--listen", action="store_true", help="Host the app on the local network"
+)
+parser.add_argument(
+    "--share", action="store_true", help="Share the app on CloudFlare tunnel"
+)
+parser.add_argument("--cpu", action="store_true", help="Run the models on the CPU")
 parser.add_argument("--summarization-model", help="Load a custom summarization model")
-parser.add_argument("--classification-model", help="Load a custom text classification model")
+parser.add_argument(
+    "--classification-model", help="Load a custom text classification model"
+)
+parser.add_argument("--captioning-model", help="Load a custom captioning model")
+parser.add_argument(
+    "--keyphrase-model", help="Load a custom keyphrase extraction model"
+)
+parser.add_argument("--prompt-model", help="Load a custom prompt generation model")
+parser.add_argument("--embedding-model", help="Load a custom text embedding model")
+
+parser.add_argument(
+    "--enable-modules",
+    action=SplitArgs,
+    default=[],
+    help="Override a list of enabled modules",
+)
 
 args = parser.parse_args()
 
+
 summarization_model = (
     args.summarization_model
     if args.summarization_model