Update server.py

server.py

@@ -1,12 +1,16 @@
+from functools import wraps
 from flask import (
     Flask,
     jsonify,
     request,
     render_template_string,
     abort,
+    send_from_directory,
+    send_file,
 )
 from flask_cors import CORS
 import unicodedata
+import argparse
 import markdown
 import time
 import os
@@ -33,11 +37,30 @@ colorama_init()
 port = 7860
 host = "0.0.0.0"
 
+
+
+class SplitArgs(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        setattr(
+            namespace, self.dest, values.replace('"', "").replace("'", "").split(",")
+        )
+parser = argparse.ArgumentParser(
+    prog="TavernAI Extras", description="Web API for transformers models"
+)
+parser.add_argument("--summarization-model", help="Load a custom summarization model")
+parser.add_argument("--classification-model", help="Load a custom text classification model")
+
+args = parser.parse_args()
+
 summarization_model = (
-    Qiliang/bart-large-cnn-samsum-ChatGPT_v3
+    args.summarization_model
+    if args.summarization_model
+    else DEFAULT_SUMMARIZATION_MODEL
 )
 classification_model = (
-    joeddav/distilbert-base-uncased-go-emotions-student
+    args.classification_model
+    if args.classification_model
+    else DEFAULT_CLASSIFICATION_MODEL
 )