Spaces:
Build error
Build error
File size: 4,498 Bytes
452b173 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from model import ExLlama, ExLlamaConfig
from flask import Flask, render_template, request, jsonify
from flask import Response, stream_with_context
from threading import Timer, Lock
import webbrowser
import json
import model_init
from session import prepare_sessions, get_initial_session, Session, load_session, new_session, _sessions_dir
import argparse
from tokenizer import ExLlamaTokenizer
from waitress import serve
app = Flask(__name__)
app.static_folder = 'static'
generate_lock = Lock()
session: Session
# Render template
@app.route("/")
def home():
return render_template("index.html")
# Get existing sessions
@app.route("/api/populate")
def api_populate():
global session
return session.api_populate()
# Edit block
@app.route("/api/edit_block", methods=['POST'])
def api_edit_block():
global session
data = request.get_json()
session.api_edit_block(data)
return json.dumps({"result": "ok"}) + "\n"
# Delete block
@app.route("/api/delete_block", methods=['POST'])
def api_delete_block():
global session
data = request.get_json()
session.api_delete_block(data)
return json.dumps({"result": "ok"}) + "\n"
# Rename session
@app.route("/api/rename_session", methods=['POST'])
def api_rename_session():
global session
data = request.get_json()
success = session.api_rename_session(data)
return json.dumps({"result": "ok" if success else "fail"}) + "\n"
# Delete session
@app.route("/api/delete_session", methods=['POST'])
def api_delete_session():
global session
data = request.get_json()
session.api_delete_session(data)
return json.dumps({"result": "ok"}) + "\n"
# Set fixed prompt settings
@app.route("/api/set_fixed_prompt", methods=['POST'])
def api_set_fixed_prompt():
global session
data = request.get_json()
session.api_set_fixed_prompt(data)
return json.dumps({"result": "ok"}) + "\n"
# Set generation settings
@app.route("/api/set_gen_settings", methods=['POST'])
def api_set_gen_settings():
global session
data = request.get_json()
session.api_set_gen_settings(data)
return json.dumps({"result": "ok"}) + "\n"
# Set session
@app.route("/api/set_session", methods=['POST'])
def api_set_session():
global session
data = request.get_json()
load_session_name = data["session_name"]
if load_session_name == ".":
session = new_session()
else:
session = load_session(load_session_name, append_path = True)
return json.dumps({"result": "ok"}) + "\n"
# Set participants
@app.route("/api/set_participants", methods=['POST'])
def api_set_participants():
global session
data = request.get_json()
session.api_set_participants(data)
return json.dumps({"result": "ok"}) + "\n"
# Accept input
@app.route("/api/userinput", methods=['POST'])
def api_userinput():
data = request.get_json()
user_input = data["user_input"]
with generate_lock:
result = Response(stream_with_context(session.respond_multi(user_input)), mimetype = 'application/json')
return result
@app.route("/api/append_block", methods=['POST'])
def api_append_block():
data = request.get_json()
session.api_append_block(data)
return json.dumps({"result": "ok"}) + "\n"
# Load the model
parser = argparse.ArgumentParser(description="Simple web-based chatbot for ExLlama")
parser.add_argument("-host", "--host", type = str, help = "IP:PORT eg, 0.0.0.0:7862", default = "localhost:5000")
parser.add_argument("-sd", "--sessions_dir", type = str, help = "Location for storing user sessions, default: ~/exllama_sessions/", default = "~/exllama_sessions/")
model_init.add_args(parser)
args = parser.parse_args()
model_init.post_parse(args)
model_init.get_model_files(args)
model_init.print_options(args)
config = model_init.make_config(args)
model_init.set_globals(args)
print(f" -- Loading model...")
model = ExLlama(config)
print(f" -- Loading tokenizer...")
tokenizer = ExLlamaTokenizer(args.tokenizer)
model_init.print_stats(model)
# Get the session ready
prepare_sessions(model, tokenizer, args.sessions_dir)
session = get_initial_session()
print(f" -- Sessions stored in: {_sessions_dir()}")
# Start the web server
machine = args.host
host, port = machine.split(":")
if host == "localhost":
Timer(1, lambda: webbrowser.open(f'http://{machine}/')).start()
serve(app, host = host, port = port) |