Upload 7 files

Dockerfile

@@ -0,0 +1,17 @@
+# Step 1: Choose a base image
+FROM python:3.9
+
+# Step 2: Set the working directory inside the container
+WORKDIR /app
+
+# Step 3: Copy the application files to the container
+COPY . /app
+
+# Step 4: Install the required dependencies
+RUN pip install -r requirements.txt
+
+# Step 5: Expose the necessary port
+EXPOSE 5000
+
+# Step 6: Define the entry point command
+CMD ["flask", "run", "--host=0.0.0.0", "--port=5000"]

app.py

@@ -0,0 +1,46 @@
+import hivemind
+from flask import Flask
+from flask_cors import CORS
+from flask_sock import Sock
+from transformers import AutoTokenizer
+
+from petals import AutoDistributedModelForCausalLM
+
+import config
+
+
+logger = hivemind.get_logger(__file__)
+
+models = {}
+for model_info in config.MODELS:
+    logger.info(f"Loading tokenizer for {model_info.repo}")
+    tokenizer = AutoTokenizer.from_pretrained(model_info.repo, add_bos_token=False, use_fast=False)
+
+    logger.info(f"Loading model {model_info.repo} with adapter {model_info.adapter} and dtype {config.TORCH_DTYPE}")
+    # We set use_fast=False since LlamaTokenizerFast takes a long time to init
+    model = AutoDistributedModelForCausalLM.from_pretrained(
+        model_info.repo,
+        active_adapter=model_info.adapter,
+        torch_dtype=config.TORCH_DTYPE,
+        initial_peers=config.INITIAL_PEERS,
+        max_retries=3,
+    )
+    model = model.to(config.DEVICE)
+
+    model_name = model_info.adapter if model_info.adapter is not None else model_info.repo
+    models[model_name] = model, tokenizer
+
+logger.info("Starting Flask app")
+app = Flask(__name__)
+CORS(app)
+app.config['SOCK_SERVER_OPTIONS'] = {'ping_interval': 25}
+sock = Sock(app)
+
+
+@app.route("/")
+def main_page():
+    return app.send_static_file("index.html")
+
+
+import http_api
+import websocket_api

config.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from cpufeature import CPUFeature
+from petals.constants import PUBLIC_INITIAL_PEERS
+
+
+@dataclass
+class ModelInfo:
+    repo: str
+    adapter: Optional[str] = None
+
+
+MODELS = [
+    ModelInfo(repo="meta-llama/Llama-2-70b-hf"),
+    ModelInfo(repo="meta-llama/Llama-2-70b-chat-hf"),
+    ModelInfo(repo="enoch/llama-65b-hf"),
+    ModelInfo(repo="enoch/llama-65b-hf", adapter="timdettmers/guanaco-65b"),
+    # ModelInfo(repo="bigscience/bloom"),
+    ModelInfo(repo="bigscience/bloomz"),
+]
+DEFAULT_MODEL_NAME = "enoch/llama-65b-hf"
+
+INITIAL_PEERS = PUBLIC_INITIAL_PEERS
+# Set this to a list of multiaddrs to connect to a private swarm instead of the public one, for example:
+# INITIAL_PEERS = ['/ip4/10.1.2.3/tcp/31234/p2p/QmcXhze98AcgGQDDYna23s4Jho96n8wkwLJv78vxtFNq44']
+
+DEVICE = "cpu"
+
+if DEVICE == "cuda":
+    TORCH_DTYPE = "auto"
+elif CPUFeature["AVX512f"] and CPUFeature["OS_AVX512"]:
+    TORCH_DTYPE = torch.bfloat16
+else:
+    TORCH_DTYPE = torch.float32  # You can use bfloat16 in this case too, but it will be slow
+
+STEP_TIMEOUT = 5 * 60
+MAX_SESSIONS = 50  # Has effect only for API v1 (HTTP-based)

http_api.py

@@ -0,0 +1,59 @@
+from traceback import format_exc
+
+import hivemind
+from flask import jsonify, request
+
+import config
+from app import app, models
+from utils import safe_decode
+
+logger = hivemind.get_logger(__file__)
+
+
+@app.post("/api/v1/generate")
+def http_api_generate():
+    try:
+        model_name = get_typed_arg("model", str, config.DEFAULT_MODEL_NAME)
+        inputs = request.values.get("inputs")
+        do_sample = get_typed_arg("do_sample", int, 0)
+        temperature = get_typed_arg("temperature", float, 1.0)
+        top_k = get_typed_arg("top_k", int)
+        top_p = get_typed_arg("top_p", float)
+        max_length = get_typed_arg("max_length", int)
+        max_new_tokens = get_typed_arg("max_new_tokens", int)
+        session_id = request.values.get("session_id")
+        logger.info(f"generate(), model={repr(model_name)}, inputs={repr(inputs)}")
+
+        if session_id is not None:
+            raise RuntimeError(
+                "Reusing inference sessions was removed from HTTP API, please use WebSocket API instead"
+            )
+
+        model, tokenizer = models[model_name]
+
+        if inputs is not None:
+            inputs = tokenizer(inputs, return_tensors="pt")["input_ids"].to(config.DEVICE)
+            n_input_tokens = inputs.shape[1]
+        else:
+            n_input_tokens = 0
+
+        outputs = model.generate(
+            inputs=inputs,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            max_length=max_length,
+            max_new_tokens=max_new_tokens,
+        )
+        outputs = safe_decode(tokenizer, outputs[0, n_input_tokens:])
+        logger.info(f"generate(), outputs={repr(outputs)}")
+
+        return jsonify(ok=True, outputs=outputs)
+    except Exception:
+        return jsonify(ok=False, traceback=format_exc())
+
+
+def get_typed_arg(name, expected_type, default=None):
+    value = request.values.get(name)
+    return expected_type(value) if value is not None else default

requirements.txt

@@ -0,0 +1,5 @@
+git+https://github.com/bigscience-workshop/petals
+Flask
+flask-sock
+flask-cors
+gunicorn[gthread]

utils.py

@@ -0,0 +1,11 @@
+import torch
+from transformers import PreTrainedTokenizerBase
+
+
+def safe_decode(tokenizer: PreTrainedTokenizerBase, outputs: torch.Tensor):
+    # Workaround to make SentencePiece .decode() keep leading spaces in a token
+    fake_token = tokenizer("^")["input_ids"][0]
+    result = tokenizer.decode([fake_token] + outputs.tolist())
+
+    # We use .lstrip() since SentencePiece may add leading spaces, e.g. if the outputs are "</s>"
+    return result.lstrip()[1:]

websocket_api.py

@@ -0,0 +1,90 @@
+import json
+from traceback import format_exc
+
+import flask_sock
+import hivemind
+import torch
+
+import config
+from app import sock, models
+from utils import safe_decode
+
+logger = hivemind.get_logger(__file__)
+
+
+@sock.route("/api/v2/generate")
+def ws_api_generate(ws):
+    try:
+        request = json.loads(ws.receive(timeout=config.STEP_TIMEOUT))
+        assert request["type"] == "open_inference_session"
+        model_name = request.get("model")
+        if model_name is None:
+            model_name = config.DEFAULT_MODEL_NAME
+        logger.info(f"ws.generate.open(), model={repr(model_name)}, max_length={repr(request['max_length'])}")
+
+        model, tokenizer = models[model_name]
+
+        with model.inference_session(max_length=request["max_length"]) as session:
+            ws.send(json.dumps({"ok": True}))
+
+            while True:
+                request = json.loads(ws.receive(timeout=config.STEP_TIMEOUT))
+                assert request["type"] == "generate"
+                inputs = request.get("inputs")
+                logger.info(f"ws.generate.step(), inputs={repr(inputs)}")
+
+                if inputs is not None:
+                    inputs = tokenizer(inputs, return_tensors="pt")["input_ids"].to(config.DEVICE)
+                    n_input_tokens = inputs.shape[1]
+                else:
+                    n_input_tokens = 0
+
+                stop_sequence = request.get("stop_sequence")
+                extra_stop_sequences = request.get("extra_stop_sequences")
+                if extra_stop_sequences is not None:
+                    cont_token = tokenizer(stop_sequence, return_tensors="pt")["input_ids"].to(config.DEVICE)
+                    assert cont_token.shape == (1, 1), \
+                        "extra_stop_sequences require stop_sequence length to be exactly 1 token"
+
+                all_outputs = ''
+                delta_q = []
+                stop = False
+                while not stop:
+                    outputs = model.generate(
+                        inputs=inputs,
+                        do_sample=request.get("do_sample", False),
+                        temperature=request.get("temperature", 1.0),
+                        top_k=request.get("top_k"),
+                        top_p=request.get("top_p"),
+                        max_length=request.get("max_length"),
+                        max_new_tokens=request.get("max_new_tokens"),
+                        session=session,
+                    )
+                    delta = outputs[0, n_input_tokens:].tolist()
+                    outputs = safe_decode(tokenizer, torch.Tensor(delta_q + delta))
+                    inputs = None  # Inputs are passed only for the 1st token of the bot's response
+                    n_input_tokens = 0 
+                    combined = all_outputs + outputs
+                    stop = stop_sequence is None or combined.endswith(stop_sequence)
+                    if extra_stop_sequences is not None:
+                        for seq in extra_stop_sequences:
+                            if combined.endswith(seq):
+                                stop = True
+                                session.last_token_id = cont_token
+                    if not stop and outputs[-10:].find(u'\ufffd') > -1:
+                        # If there's a replacement character, keep getting more tokens
+                        # until we can decode properly
+                        delta_q = delta_q + delta
+                        logger.info(f"ws.generate.append_retry(), all_outputs={repr(combined)}")
+                    else:
+                        all_outputs = combined
+                        delta_q = []
+                        logger.info(f"ws.generate.step(), all_outputs={repr(all_outputs)}, stop={stop}")
+                        ws.send(json.dumps({"ok": True, "outputs": outputs, "stop": stop}))
+    except flask_sock.ConnectionClosed:
+        pass
+    except Exception:
+        logger.warning("ws.generate failed:", exc_info=True)
+        ws.send(json.dumps({"ok": False, "traceback": format_exc()}))
+    finally:
+        logger.info(f"ws.generate.close()")