Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 5,080 Bytes
6f40009 30c8aac 6f40009 30c8aac 6f40009 30c8aac 6f40009 6d90da5 6f40009 6d90da5 6f40009 30c8aac |
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 |
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
from urllib.parse import urlparse, unquote
model: RobertaForSequenceClassification = None
tokenizer: RobertaTokenizer = None
device: str = None
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 = unquote(urlparse(self.path).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)
# avoid calling CUDA API before forking; doing so in a subprocess is fine.
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)
|