Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
Remove useless file
Browse files- detector/server_get.py +0 -120
detector/server_get.py
DELETED
@@ -1,120 +0,0 @@
|
|
1 |
-
import os
|
2 |
-
import sys
|
3 |
-
from http.server import HTTPServer, SimpleHTTPRequestHandler
|
4 |
-
from multiprocessing import Process
|
5 |
-
import subprocess
|
6 |
-
from transformers import RobertaForSequenceClassification, RobertaTokenizer
|
7 |
-
import json
|
8 |
-
import fire
|
9 |
-
import torch
|
10 |
-
from urllib.parse import urlparse, unquote
|
11 |
-
|
12 |
-
|
13 |
-
model: RobertaForSequenceClassification = None
|
14 |
-
tokenizer: RobertaTokenizer = None
|
15 |
-
device: str = None
|
16 |
-
|
17 |
-
def log(*args):
|
18 |
-
print(f"[{os.environ.get('RANK', '')}]", *args, file=sys.stderr)
|
19 |
-
|
20 |
-
|
21 |
-
class RequestHandler(SimpleHTTPRequestHandler):
|
22 |
-
|
23 |
-
def do_GET(self):
|
24 |
-
query = unquote(urlparse(self.path).query)
|
25 |
-
|
26 |
-
if not query:
|
27 |
-
self.begin_content('text/html')
|
28 |
-
|
29 |
-
html = os.path.join(os.path.dirname(__file__), 'index.html')
|
30 |
-
self.wfile.write(open(html).read().encode())
|
31 |
-
return
|
32 |
-
|
33 |
-
self.begin_content('application/json;charset=UTF-8')
|
34 |
-
|
35 |
-
tokens = tokenizer.encode(query)
|
36 |
-
all_tokens = len(tokens)
|
37 |
-
tokens = tokens[:tokenizer.max_len - 2]
|
38 |
-
used_tokens = len(tokens)
|
39 |
-
tokens = torch.tensor([tokenizer.bos_token_id] + tokens + [tokenizer.eos_token_id]).unsqueeze(0)
|
40 |
-
mask = torch.ones_like(tokens)
|
41 |
-
|
42 |
-
with torch.no_grad():
|
43 |
-
logits = model(tokens.to(device), attention_mask=mask.to(device))[0]
|
44 |
-
probs = logits.softmax(dim=-1)
|
45 |
-
|
46 |
-
fake, real = probs.detach().cpu().flatten().numpy().tolist()
|
47 |
-
|
48 |
-
self.wfile.write(json.dumps(dict(
|
49 |
-
all_tokens=all_tokens,
|
50 |
-
used_tokens=used_tokens,
|
51 |
-
real_probability=real,
|
52 |
-
fake_probability=fake
|
53 |
-
)).encode())
|
54 |
-
|
55 |
-
def begin_content(self, content_type):
|
56 |
-
self.send_response(200)
|
57 |
-
self.send_header('Content-Type', content_type)
|
58 |
-
self.send_header('Access-Control-Allow-Origin', '*')
|
59 |
-
self.end_headers()
|
60 |
-
|
61 |
-
def log_message(self, format, *args):
|
62 |
-
log(format % args)
|
63 |
-
|
64 |
-
|
65 |
-
def serve_forever(server, model, tokenizer, device):
|
66 |
-
log('Process has started; loading the model ...')
|
67 |
-
globals()['model'] = model.to(device)
|
68 |
-
globals()['tokenizer'] = tokenizer
|
69 |
-
globals()['device'] = device
|
70 |
-
|
71 |
-
log(f'Ready to serve at http://localhost:{server.server_address[1]}')
|
72 |
-
server.serve_forever()
|
73 |
-
|
74 |
-
|
75 |
-
def main(checkpoint, port=8080, device='cuda' if torch.cuda.is_available() else 'cpu'):
|
76 |
-
if checkpoint.startswith('gs://'):
|
77 |
-
print(f'Downloading {checkpoint}', file=sys.stderr)
|
78 |
-
subprocess.check_output(['gsutil', 'cp', checkpoint, '.'])
|
79 |
-
checkpoint = os.path.basename(checkpoint)
|
80 |
-
assert os.path.isfile(checkpoint)
|
81 |
-
|
82 |
-
print(f'Loading checkpoint from {checkpoint}')
|
83 |
-
data = torch.load(checkpoint, map_location='cpu')
|
84 |
-
|
85 |
-
model_name = 'roberta-large' if data['args']['large'] else 'roberta-base'
|
86 |
-
model = RobertaForSequenceClassification.from_pretrained(model_name)
|
87 |
-
tokenizer = RobertaTokenizer.from_pretrained(model_name)
|
88 |
-
|
89 |
-
model.load_state_dict(data['model_state_dict'])
|
90 |
-
model.eval()
|
91 |
-
|
92 |
-
print(f'Starting HTTP server on port {port}', file=sys.stderr)
|
93 |
-
server = HTTPServer(('0.0.0.0', port), RequestHandler)
|
94 |
-
|
95 |
-
# avoid calling CUDA API before forking; doing so in a subprocess is fine.
|
96 |
-
num_workers = int(subprocess.check_output([sys.executable, '-c', 'import torch; print(torch.cuda.device_count())']))
|
97 |
-
|
98 |
-
if num_workers <= 1:
|
99 |
-
serve_forever(server, model, tokenizer, device)
|
100 |
-
else:
|
101 |
-
print(f'Launching {num_workers} worker processes...')
|
102 |
-
|
103 |
-
subprocesses = []
|
104 |
-
|
105 |
-
for i in range(num_workers):
|
106 |
-
os.environ['RANK'] = f'{i}'
|
107 |
-
os.environ['CUDA_VISIBLE_DEVICES'] = f'{i}'
|
108 |
-
process = Process(target=serve_forever, args=(server, model, tokenizer, device))
|
109 |
-
process.start()
|
110 |
-
subprocesses.append(process)
|
111 |
-
|
112 |
-
del os.environ['RANK']
|
113 |
-
del os.environ['CUDA_VISIBLE_DEVICES']
|
114 |
-
|
115 |
-
for process in subprocesses:
|
116 |
-
process.join()
|
117 |
-
|
118 |
-
|
119 |
-
if __name__ == '__main__':
|
120 |
-
fire.Fire(main)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|