XciD HF staff commited on
Commit
999a35e
1 Parent(s): 4c9574e

Remove useless file

Browse files
Files changed (1) hide show
  1. 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)