refactor: remove unused python code, move trainer to python/
Meng Zhang
parent
0f72788d82
commit
2691b302f0
|
|
@ -1,69 +0,0 @@
|
|||
import os
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, Response
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from . import events as events_lib
|
||||
from .backend import PythonModelService, TritonService
|
||||
from .models import (
|
||||
ChoiceEvent,
|
||||
CompletionEvent,
|
||||
CompletionRequest,
|
||||
CompletionResponse,
|
||||
EventTypeMapping,
|
||||
)
|
||||
|
||||
app = FastAPI(
|
||||
title="TabbyServer",
|
||||
description="""
|
||||
[](http://github.com/TabbyML/tabby)
|
||||
|
||||
TabbyServer is the backend for tabby, serving code completion requests from code editor / IDE.
|
||||
|
||||
* [Admin Panel](./_admin)
|
||||
""",
|
||||
docs_url="/",
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=False,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
MODEL_NAME = os.environ.get("MODEL_NAME")
|
||||
MODEL_BACKEND = os.environ.get("MODEL_BACKEND", "python")
|
||||
|
||||
if MODEL_BACKEND == "triton":
|
||||
model_backend = TritonService(
|
||||
tokenizer_name=MODEL_NAME,
|
||||
host=os.environ.get("TRITON_HOST", "localhost"),
|
||||
port=os.environ.get("TRITON_PORT", "8001"),
|
||||
)
|
||||
else:
|
||||
model_backend = PythonModelService(MODEL_NAME)
|
||||
|
||||
LOGS_DIR = os.environ.get("LOGS_DIR", None)
|
||||
if LOGS_DIR is not None:
|
||||
events_lib.setup_logging(os.path.join(LOGS_DIR, "tabby-server"))
|
||||
|
||||
|
||||
@app.post("/v1/completions")
|
||||
async def completions(request: CompletionRequest) -> CompletionResponse:
|
||||
response = model_backend(request)
|
||||
events_lib.log_completion(request, response)
|
||||
return response
|
||||
|
||||
|
||||
@app.post("/v1/events")
|
||||
async def events(e: ChoiceEvent | CompletionEvent) -> JSONResponse:
|
||||
if isinstance(e, EventTypeMapping[e.type]):
|
||||
events_lib.log_event(e)
|
||||
return JSONResponse(content="ok")
|
||||
else:
|
||||
print(type(e))
|
||||
return JSONResponse(content="invalid event", status_code=422)
|
||||
|
|
@ -1,2 +0,0 @@
|
|||
from .python import PythonModelService
|
||||
from .triton import TritonService
|
||||
|
|
@ -1,77 +0,0 @@
|
|||
from typing import List, Optional, Set
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..models import Language
|
||||
|
||||
|
||||
class LanguagePreset(BaseModel):
|
||||
max_length: int
|
||||
stop_words: List[str]
|
||||
reserved_keywords: Optional[Set]
|
||||
|
||||
|
||||
LanguagePresets = {
|
||||
Language.UNKNOWN: LanguagePreset(
|
||||
max_length=128,
|
||||
stop_words=["\n\n"],
|
||||
),
|
||||
Language.PYTHON: LanguagePreset(
|
||||
max_length=128,
|
||||
stop_words=["\n\n", "\ndef", "\n#", "\nimport", "\nfrom", "\nclass"],
|
||||
reserved_keywords=set(
|
||||
[
|
||||
"False",
|
||||
"class",
|
||||
"from",
|
||||
"or",
|
||||
"None",
|
||||
"continue",
|
||||
"global",
|
||||
"pass",
|
||||
"True",
|
||||
"def",
|
||||
"if",
|
||||
"raise",
|
||||
"and",
|
||||
"del",
|
||||
"import",
|
||||
"return",
|
||||
"as",
|
||||
"elif",
|
||||
"in",
|
||||
"try",
|
||||
"assert",
|
||||
"else",
|
||||
"is",
|
||||
"while",
|
||||
"async",
|
||||
"except",
|
||||
"lambda",
|
||||
"with",
|
||||
"await",
|
||||
"finally",
|
||||
"nonlocal",
|
||||
"yield",
|
||||
"break",
|
||||
"for",
|
||||
"not",
|
||||
]
|
||||
),
|
||||
),
|
||||
Language.JAVASCRIPT: LanguagePreset(
|
||||
max_length=128, stop_words=["\n\n", "\nfunction", "\n//", "\nimport", "\nclass"]
|
||||
),
|
||||
Language.TYPESCRIPT: LanguagePreset(
|
||||
max_length=128,
|
||||
stop_words=[
|
||||
"\n\n",
|
||||
"\nfunction",
|
||||
"\n//",
|
||||
"\nimport",
|
||||
"\nclass",
|
||||
"\ninterface",
|
||||
"\ntype",
|
||||
],
|
||||
),
|
||||
}
|
||||
|
|
@ -1,67 +0,0 @@
|
|||
import time
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
StoppingCriteria,
|
||||
StoppingCriteriaList,
|
||||
)
|
||||
|
||||
from ..models import Choice, CompletionRequest, CompletionResponse
|
||||
from .language_presets import LanguagePresets
|
||||
from .utils import random_completion_id, trim_with_stop_words
|
||||
|
||||
|
||||
class PythonModelService:
|
||||
def __init__(self, model_name, quantize=False):
|
||||
if torch.cuda.is_available():
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
|
||||
else:
|
||||
if quantize:
|
||||
raise ValueError("quantization on CPU is not implemented")
|
||||
|
||||
device = torch.device("cpu")
|
||||
dtype = torch.float32
|
||||
|
||||
self.device = device
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_name, local_files_only=True
|
||||
)
|
||||
self.model = (
|
||||
AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
torch_dtype=dtype,
|
||||
device_map="auto" if torch.cuda.is_available() else None,
|
||||
local_files_only=True,
|
||||
load_in_8bit=quantize,
|
||||
)
|
||||
.to(device)
|
||||
.eval()
|
||||
)
|
||||
|
||||
def generate(self, request: CompletionRequest) -> List[Choice]:
|
||||
preset = LanguagePresets.get(request.language, None)
|
||||
if preset is None:
|
||||
return []
|
||||
|
||||
input_ids = self.tokenizer.encode(request.prompt, return_tensors="pt").to(
|
||||
self.device
|
||||
)
|
||||
res = self.model.generate(
|
||||
input_ids,
|
||||
max_length=preset.max_length,
|
||||
)
|
||||
output_ids = res[0][len(input_ids[0]) :]
|
||||
text = trim_with_stop_words(
|
||||
self.tokenizer.decode(output_ids), preset.stop_words
|
||||
)
|
||||
return [Choice(index=0, text=text)]
|
||||
|
||||
def __call__(self, request: CompletionRequest) -> CompletionResponse:
|
||||
choices = self.generate(request)
|
||||
return CompletionResponse(
|
||||
id=random_completion_id(), created=int(time.time()), choices=choices
|
||||
)
|
||||
|
|
@ -1,121 +0,0 @@
|
|||
import time
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import tritonclient.grpc as client_util
|
||||
from transformers import AutoTokenizer
|
||||
from tritonclient.utils import InferenceServerException, np_to_triton_dtype
|
||||
|
||||
from ..models import Choice, CompletionRequest, CompletionResponse
|
||||
from .language_presets import LanguagePresets
|
||||
from .utils import random_completion_id, trim_with_stop_words
|
||||
|
||||
|
||||
class TritonService:
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer_name,
|
||||
host: str = "localhost",
|
||||
port: int = 8001,
|
||||
verbose: bool = False,
|
||||
):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
|
||||
self.client = client_util.InferenceServerClient(
|
||||
url=f"{host}:{port}", verbose=verbose
|
||||
)
|
||||
|
||||
def generate(self, data: CompletionRequest) -> List[Choice]:
|
||||
n = 1
|
||||
np_type = np.uint32
|
||||
model_name = "fastertransformer"
|
||||
|
||||
preset = LanguagePresets.get(data.language, None)
|
||||
if preset is None:
|
||||
return []
|
||||
|
||||
prompt = data.prompt
|
||||
|
||||
input_start_ids = np.expand_dims(self.tokenizer.encode(prompt), 0)
|
||||
input_start_ids = np.repeat(input_start_ids, n, axis=0).astype(np_type)
|
||||
prompt_len = input_start_ids.shape[1]
|
||||
input_len = prompt_len * np.ones([input_start_ids.shape[0], 1]).astype(np_type)
|
||||
|
||||
prompt_tokens: int = input_len[0][0]
|
||||
output_len = np.ones_like(input_len).astype(np_type) * preset.max_length
|
||||
|
||||
stop_word_list = np.repeat(
|
||||
to_word_list_format([preset.stop_words], self.tokenizer),
|
||||
input_start_ids.shape[0],
|
||||
axis=0,
|
||||
)
|
||||
|
||||
inputs = [
|
||||
prepare_tensor("input_ids", input_start_ids),
|
||||
prepare_tensor("input_lengths", input_len),
|
||||
prepare_tensor("request_output_len", output_len),
|
||||
prepare_tensor("stop_words_list", stop_word_list),
|
||||
]
|
||||
|
||||
result = self.client.infer(model_name, inputs)
|
||||
|
||||
output_data = result.as_numpy("output_ids")
|
||||
if output_data is None:
|
||||
raise RuntimeError("No output data")
|
||||
|
||||
output_data = output_data.squeeze(1)
|
||||
sequence_lengths = result.as_numpy("sequence_length").squeeze(1)
|
||||
gen_len = sequence_lengths - input_len.squeeze(1)
|
||||
|
||||
decoded = [
|
||||
self.tokenizer.decode(out[prompt_len : prompt_len + g])
|
||||
for g, out in zip(gen_len, output_data)
|
||||
]
|
||||
trimmed = [trim_with_stop_words(d, preset.stop_words) for d in decoded]
|
||||
return [Choice(index=i, text=text) for i, text in enumerate(trimmed)]
|
||||
|
||||
def __call__(self, data: CompletionRequest) -> CompletionResponse:
|
||||
choices = self.generate(data)
|
||||
return CompletionResponse(
|
||||
id=random_completion_id(), created=int(time.time()), choices=choices
|
||||
)
|
||||
|
||||
|
||||
def prepare_tensor(name: str, tensor_input):
|
||||
t = client_util.InferInput(
|
||||
name, tensor_input.shape, np_to_triton_dtype(tensor_input.dtype)
|
||||
)
|
||||
t.set_data_from_numpy(tensor_input)
|
||||
return t
|
||||
|
||||
|
||||
def to_word_list_format(word_dict, tokenizer):
|
||||
flat_ids = []
|
||||
offsets = []
|
||||
for word_dict_item in word_dict:
|
||||
item_flat_ids = []
|
||||
item_offsets = []
|
||||
|
||||
for word in word_dict_item:
|
||||
ids = tokenizer.encode(word)
|
||||
|
||||
if len(ids) == 0:
|
||||
continue
|
||||
|
||||
item_flat_ids += ids
|
||||
item_offsets.append(len(ids))
|
||||
|
||||
if word == "\n\n":
|
||||
ids = tokenizer.encode("\n") * 2
|
||||
item_flat_ids += ids
|
||||
item_offsets.append(len(ids))
|
||||
|
||||
flat_ids.append(np.array(item_flat_ids))
|
||||
offsets.append(np.cumsum(np.array(item_offsets)))
|
||||
|
||||
pad_to = max(1, max(len(ids) for ids in flat_ids))
|
||||
|
||||
for i, (ids, offs) in enumerate(zip(flat_ids, offsets)):
|
||||
flat_ids[i] = np.pad(ids, (0, pad_to - len(ids)), constant_values=0)
|
||||
offsets[i] = np.pad(offs, (0, pad_to - len(offs)), constant_values=-1)
|
||||
|
||||
return np.array([flat_ids, offsets], dtype=np.int32).transpose((1, 0, 2))
|
||||
|
|
@ -1,16 +0,0 @@
|
|||
import random
|
||||
import string
|
||||
|
||||
|
||||
def random_completion_id():
|
||||
return "cmpl-" + "".join(
|
||||
random.choice(string.ascii_letters + string.digits) for _ in range(29)
|
||||
)
|
||||
|
||||
|
||||
def trim_with_stop_words(output: str, stopwords: list) -> str:
|
||||
for w in sorted(stopwords, key=len, reverse=True):
|
||||
index = output.find(w)
|
||||
if index != -1:
|
||||
output = output[:index]
|
||||
return output
|
||||
|
|
@ -1,37 +0,0 @@
|
|||
import os
|
||||
import shutil
|
||||
|
||||
from loguru import logger
|
||||
from pydantic import BaseModel
|
||||
|
||||
from . import models
|
||||
|
||||
|
||||
def setup_logging(logdir):
|
||||
try:
|
||||
shutil.rmtree(logdir)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
# Remove default handler
|
||||
logger.add(
|
||||
os.path.join(logdir, "events.{time}.log"),
|
||||
rotation="1 hours",
|
||||
retention="2 hours",
|
||||
level="INFO",
|
||||
filter=__name__,
|
||||
enqueue=True,
|
||||
delay=True,
|
||||
serialize=True,
|
||||
)
|
||||
|
||||
|
||||
def log_completion(
|
||||
request: models.CompletionRequest, response: models.CompletionResponse
|
||||
) -> None:
|
||||
event = models.CompletionEvent.build(request, response)
|
||||
logger.info(event.json())
|
||||
|
||||
|
||||
def log_event(event: models.Event):
|
||||
logger.info(event.json())
|
||||
|
|
@ -1,77 +0,0 @@
|
|||
from enum import Enum
|
||||
from typing import List
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class Choice(BaseModel):
|
||||
index: int
|
||||
text: str
|
||||
|
||||
|
||||
# https://code.visualstudio.com/docs/languages/identifiers
|
||||
class Language(str, Enum):
|
||||
UNKNOWN = "unknown"
|
||||
PYTHON = "python"
|
||||
JAVASCRIPT = "javascript"
|
||||
TYPESCRIPT = "typescript"
|
||||
|
||||
|
||||
class CompletionRequest(BaseModel):
|
||||
language: str = Field(
|
||||
example=Language.PYTHON,
|
||||
default=Language.UNKNOWN,
|
||||
description="Language for completion request",
|
||||
)
|
||||
|
||||
prompt: str = Field(
|
||||
example="def binarySearch(arr, left, right, x):\n mid = (left +",
|
||||
description="The context to generate completions for, encoded as a string.",
|
||||
)
|
||||
|
||||
|
||||
class CompletionResponse(BaseModel):
|
||||
id: str
|
||||
created: int
|
||||
choices: List[Choice]
|
||||
|
||||
|
||||
class EventType(str, Enum):
|
||||
COMPLETION = "completion"
|
||||
VIEW = "view"
|
||||
SELECT = "select"
|
||||
|
||||
|
||||
class Event(BaseModel):
|
||||
type: EventType
|
||||
|
||||
|
||||
class CompletionEvent(Event):
|
||||
id: str
|
||||
language: str
|
||||
prompt: str
|
||||
created: int
|
||||
choices: List[Choice]
|
||||
|
||||
@classmethod
|
||||
def build(cls, request: CompletionRequest, response: CompletionResponse):
|
||||
return cls(
|
||||
type=EventType.COMPLETION,
|
||||
id=response.id,
|
||||
language=request.language,
|
||||
prompt=request.prompt,
|
||||
created=response.created,
|
||||
choices=response.choices,
|
||||
)
|
||||
|
||||
|
||||
class ChoiceEvent(Event):
|
||||
completion_id: str
|
||||
choice_index: int
|
||||
|
||||
|
||||
EventTypeMapping = {
|
||||
EventType.COMPLETION: CompletionEvent,
|
||||
EventType.VIEW: ChoiceEvent,
|
||||
EventType.SELECT: ChoiceEvent,
|
||||
}
|
||||
|
|
@ -1,11 +0,0 @@
|
|||
schedule: "*/3 * * * *" # Run every 3rd minute
|
||||
|
||||
env:
|
||||
- PATH: "$PATH"
|
||||
- LOGS_DIR: "$LOGS_DIR"
|
||||
- DB_FILE: "$DB_FILE"
|
||||
- APP_DIR: /home/app
|
||||
steps:
|
||||
- name: Collect Tabby
|
||||
dir: $APP_DIR
|
||||
command: ./tabby/tools/analytic/main.sh collect_tabby_server_logs
|
||||
|
|
@ -1,19 +0,0 @@
|
|||
schedule: "5 4 * * *" # Run daily at 04:05.
|
||||
|
||||
env:
|
||||
- PATH: "$PATH"
|
||||
- APP_DIR: /home/app
|
||||
- CONFIG_FILE: "$CONFIG_FILE"
|
||||
- GIT_REPOSITORIES_DIR: "$GIT_REPOSITORIES_DIR"
|
||||
- DATASET_DIR: "$DATASET_DIR"
|
||||
- HOME: "$HOME"
|
||||
steps:
|
||||
- name: update repositories
|
||||
dir: $APP_DIR
|
||||
command: python -m tabby.tools.repository.updater --data_dir=$GIT_REPOSITORIES_DIR --config_file=$CONFIG_FILE
|
||||
|
||||
- name: generate dataset
|
||||
dir: $APP_DIR
|
||||
command: python -m tabby.tools.build_dataset --project_dir=$GIT_REPOSITORIES_DIR --output_dir=$DATASET_DIR
|
||||
depends:
|
||||
- update repositories
|
||||
Loading…
Reference in New Issue