|
import os |
|
import sys |
|
from http.server import HTTPServer, SimpleHTTPRequestHandler |
|
from multiprocessing import Process |
|
import subprocess |
|
from transformers import RobertaForSequenceClassification, RobertaTokenizer |
|
import json |
|
import fire |
|
import torch |
|
import re |
|
from urllib.parse import urlparse, unquote, parse_qs, urlencode |
|
|
|
model: RobertaForSequenceClassification = None |
|
tokenizer: RobertaTokenizer = None |
|
device: str = None |
|
|
|
|
|
regex = r"__theme=(.+)" |
|
|
|
|
|
def log(*args): |
|
print(f"[{os.environ.get('RANK', '')}]", *args, file=sys.stderr) |
|
|
|
|
|
class RequestHandler(SimpleHTTPRequestHandler): |
|
|
|
def do_POST(self): |
|
self.begin_content('application/json,charset=UTF-8') |
|
|
|
content_length = int(self.headers['Content-Length']) |
|
if content_length > 0: |
|
post_data = self.rfile.read(content_length).decode('utf-8') |
|
try: |
|
post_data = json.loads(post_data) |
|
|
|
if 'text' not in post_data: |
|
self.wfile.write(json.dumps({"error": "missing key 'text'"}).encode('utf-8')) |
|
else: |
|
all_tokens, used_tokens, fake, real = self.infer(post_data['text']) |
|
|
|
self.wfile.write(json.dumps(dict( |
|
all_tokens=all_tokens, |
|
used_tokens=used_tokens, |
|
real_probability=real, |
|
fake_probability=fake |
|
)).encode('utf-8')) |
|
|
|
except Exception as e: |
|
self.wfile.write(json.dumps({"error": str(e)}).encode('utf-8')) |
|
|
|
def do_GET(self): |
|
query = urlparse(self.path).query |
|
query = re.sub(regex, "", query, 0, re.MULTILINE) |
|
query = unquote(query) |
|
|
|
if not query: |
|
self.begin_content('text/html') |
|
|
|
html = os.path.join(os.path.dirname(__file__), 'index.html') |
|
self.wfile.write(open(html).read().encode()) |
|
return |
|
|
|
self.begin_content('application/json;charset=UTF-8') |
|
|
|
all_tokens, used_tokens, fake, real = self.infer(query) |
|
|
|
self.wfile.write(json.dumps(dict( |
|
all_tokens=all_tokens, |
|
used_tokens=used_tokens, |
|
real_probability=real, |
|
fake_probability=fake |
|
)).encode()) |
|
|
|
def infer(self, query): |
|
tokens = tokenizer.encode(query) |
|
all_tokens = len(tokens) |
|
tokens = tokens[:tokenizer.max_len - 2] |
|
used_tokens = len(tokens) |
|
tokens = torch.tensor([tokenizer.bos_token_id] + tokens + [tokenizer.eos_token_id]).unsqueeze(0) |
|
mask = torch.ones_like(tokens) |
|
|
|
with torch.no_grad(): |
|
logits = model(tokens.to(device), attention_mask=mask.to(device))[0] |
|
probs = logits.softmax(dim=-1) |
|
|
|
fake, real = probs.detach().cpu().flatten().numpy().tolist() |
|
|
|
return all_tokens, used_tokens, fake, real |
|
|
|
def begin_content(self, content_type): |
|
self.send_response(200) |
|
self.send_header('Content-Type', content_type) |
|
self.send_header('Access-Control-Allow-Origin', '*') |
|
self.end_headers() |
|
|
|
def log_message(self, format, *args): |
|
log(format % args) |
|
|
|
|
|
def serve_forever(server, model, tokenizer, device): |
|
log('Process has started; loading the model ...') |
|
globals()['model'] = model.to(device) |
|
globals()['tokenizer'] = tokenizer |
|
globals()['device'] = device |
|
|
|
log(f'Ready to serve at http://localhost:{server.server_address[1]}') |
|
server.serve_forever() |
|
|
|
|
|
def main(checkpoint, port=8080, device='cuda' if torch.cuda.is_available() else 'cpu'): |
|
if checkpoint.startswith('gs://'): |
|
print(f'Downloading {checkpoint}', file=sys.stderr) |
|
subprocess.check_output(['gsutil', 'cp', checkpoint, '.']) |
|
checkpoint = os.path.basename(checkpoint) |
|
assert os.path.isfile(checkpoint) |
|
|
|
print(f'Loading checkpoint from {checkpoint}') |
|
data = torch.load(checkpoint, map_location='cpu') |
|
|
|
model_name = 'roberta-large' if data['args']['large'] else 'roberta-base' |
|
model = RobertaForSequenceClassification.from_pretrained(model_name) |
|
tokenizer = RobertaTokenizer.from_pretrained(model_name) |
|
|
|
model.load_state_dict(data['model_state_dict']) |
|
model.eval() |
|
|
|
print(f'Starting HTTP server on port {port}', file=sys.stderr) |
|
server = HTTPServer(('0.0.0.0', port), RequestHandler) |
|
|
|
|
|
num_workers = int(subprocess.check_output([sys.executable, '-c', 'import torch; print(torch.cuda.device_count())'])) |
|
|
|
if num_workers <= 1: |
|
serve_forever(server, model, tokenizer, device) |
|
else: |
|
print(f'Launching {num_workers} worker processes...') |
|
|
|
subprocesses = [] |
|
|
|
for i in range(num_workers): |
|
os.environ['RANK'] = f'{i}' |
|
os.environ['CUDA_VISIBLE_DEVICES'] = f'{i}' |
|
process = Process(target=serve_forever, args=(server, model, tokenizer, device)) |
|
process.start() |
|
subprocesses.append(process) |
|
|
|
del os.environ['RANK'] |
|
del os.environ['CUDA_VISIBLE_DEVICES'] |
|
|
|
for process in subprocesses: |
|
process.join() |
|
|
|
|
|
if __name__ == '__main__': |
|
fire.Fire(main) |
|
|
