diff --git a/README.md b/README.md index e6afdea..6ebd0bb 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,36 @@ # TavernAI - Extras ## What is this A set of unofficial APIs for various [TavernAI](https://github.com/TavernAI/TavernAI) extensions + ## How to run * Install Python 3.10 * Run `pip install -r requirements.txt` * Run `python server.py` + ## Included functionality -* BART model for text summarization -* BERT model for text classification -* BLIP model for image captioning \ No newline at end of file +### BLIP model for image captioning +`POST /api/caption` +#### **Input** +``` +{ "image": "base64 encoded image" } +``` +#### **Output** +``` +{ "caption": "caption of the posted image" } +``` + +### BART model for text summarization +Not implemented yet + +### BERT model for text classification +Not implemented yet + +## Additional options +| Flag | Description | +| -------------- | -------------------------------------------------------------------- | +| `--port` | Specify the port on which the application is hosted. Default: *5100* | +| `--listen` | Hosts the app on the local network | +| `--share` | Shares the app on CloudFlare tunnel | +| `--bart-model` | Load a custom BART model.
Expects a HuggingFace model ID.
Default: [Qiliang/bart-large-cnn-samsum-ChatGPT_v3](https://huggingface.co/Qiliang/bart-large-cnn-samsum-ChatGPT_v3) | +| `--bert-model` | Load a custom BERT model.
Expects a HuggingFace model ID.
Default: [bhadresh-savani/distilbert-base-uncased-emotion](https://huggingface.co/bhadresh-savani/distilbert-base-uncased-emotion) | +| `--blip-model` | Load a custom BLIP model.
Expects a HuggingFace model Id.
Default: [Salesforce/blip-image-captioning-base](https://huggingface.co/Salesforce/blip-image-captioning-base) | \ No newline at end of file diff --git a/server.py b/server.py index 5f168f7..10d99e6 100644 --- a/server.py +++ b/server.py @@ -15,18 +15,15 @@ from io import BytesIO DEFAULT_BART = 'Qiliang/bart-large-cnn-samsum-ChatGPT_v3' DEFAULT_BERT = 'bhadresh-savani/distilbert-base-uncased-emotion' DEFAULT_BLIP = 'Salesforce/blip-image-captioning-base' -UPLOAD_FOLDER = './uploads' # Script arguments -parser = argparse.ArgumentParser( - prog = 'TavernAI Extras', - description = 'Web API for transformers models') +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="Hosts the app on the local network") -parser.add_argument('--share', action='store_true', help="Shares the app on cloudflare tunnel") -parser.add_argument('--bart-model', help="Customize a BART model to be used by the app") -parser.add_argument('--bert-model', help="Customize a BERT model to be used by the app") -parser.add_argument('--blip-model', help="Customize a BLIP model to be used by the app") +parser.add_argument('--share', action='store_true', help="Shares the app on CloudFlare tunnel") +parser.add_argument('--bart-model', help="Load a custom BART model") +parser.add_argument('--bert-model', help="Load a custom BERT model") +parser.add_argument('--blip-model', help="Load a custom BLIP model") args = parser.parse_args() @@ -74,7 +71,6 @@ bert = AutoModelForSequenceClassification.from_pretrained(bert_model) # Flask init app = Flask(__name__) -app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024 # AI stuff