Spaces:
Runtime error
Runtime error
File size: 2,706 Bytes
c942159 |
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 |
from flask import Flask, render_template, redirect, request, jsonify, make_response
import datetime
import torch
import transformers
device = torch.device('cuda')
MODEL_NAME = 'liujch1998/vera'
class Interactive:
def __init__(self):
self.tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
self.model = transformers.T5EncoderModel.from_pretrained(MODEL_NAME, low_cpu_mem_usage=True, device_map='auto', torch_dtype='auto', offload_folder='offload')
self.model.D = self.model.shared.embedding_dim
self.linear = torch.nn.Linear(self.model.D, 1, dtype=self.model.dtype).to(device)
self.linear.weight = torch.nn.Parameter(self.model.shared.weight[32099, :].unsqueeze(0)) # (1, D)
self.linear.bias = torch.nn.Parameter(self.model.shared.weight[32098, 0].unsqueeze(0)) # (1)
self.model.eval()
self.t = self.model.shared.weight[32097, 0].item()
def run(self, statement):
input_ids = self.tokenizer.batch_encode_plus([statement], return_tensors='pt', padding='longest', truncation='longest_first', max_length=128).input_ids.to(device)
with torch.no_grad():
output = self.model(input_ids)
last_hidden_state = output.last_hidden_state.to(device) # (B=1, L, D)
hidden = last_hidden_state[0, -1, :] # (D)
logit = self.linear(hidden).squeeze(-1) # ()
logit_calibrated = logit / self.t
score = logit.sigmoid()
score_calibrated = logit_calibrated.sigmoid()
return {
'timestamp': datetime.datetime.now().strftime('%Y%m%d-%H%M%S'),
'statement': statement,
'logit': logit.item(),
'logit_calibrated': logit_calibrated.item(),
'score': score.item(),
'score_calibrated': score_calibrated.item(),
}
interactive = Interactive()
app = Flask(__name__)
@app.route('/', methods=['GET', 'POST'])
def main():
try:
print(request)
data = request.get_json()
statement = data.get('statement')
except Exception as e:
return jsonify({
'success': False,
'message': 'Please provide a statement.',
}), 400
try:
result = interactive.run(statement)
except Exception as e:
return jsonify({
'success': False,
'message': 'Internal error.',
}), 500
return jsonify(result)
if __name__ == "__main__":
app.run(host="0.0.0.0", port=8372, threaded=True, ssl_context=('/etc/letsencrypt/live/qa.cs.washington.edu/fullchain.pem', '/etc/letsencrypt/live/qa.cs.washington.edu/privkey.pem'))
# 8372 is when you type Vera on a phone keypad
|