johnpaulbin/beanbox-toxicity

1from flask import Flask, request, jsonify2import asyncio3from hypercorn.asyncio import serve4from hypercorn.config import Config5import torch.nn.functional as F6from torch import nn7import os8os.environ['CURL_CA_BUNDLE'] = ''9 10app = Flask(__name__)11 12 13from sentence_transformers import SentenceTransformer14sentencemodel = SentenceTransformer('johnpaulbin/toxic-gte-small-3')15 16USE_GPU = False17 18 19""" Use torchMoji to predict emojis from a single text input20"""21 22import numpy as np23import emoji, json24from torchmoji.global_variables import PRETRAINED_PATH, VOCAB_PATH25from torchmoji.sentence_tokenizer import SentenceTokenizer26from torchmoji.model_def import torchmoji_emojis27import torch28 29# Emoji map in emoji_overview.png30EMOJIS = ":joy: :unamused: :weary: :sob: :heart_eyes: \31:pensive: :ok_hand: :blush: :heart: :smirk: \32:grin: :notes: :flushed: :100: :sleeping: \33:relieved: :relaxed: :raised_hands: :two_hearts: :expressionless: \34:sweat_smile: :pray: :confused: :kissing_heart: :heartbeat: \35:neutral_face: :information_desk_person: :disappointed: :see_no_evil: :tired_face: \36:v: :sunglasses: :rage: :thumbsup: :cry: \37:sleepy: :yum: :triumph: :hand: :mask: \38:clap: :eyes: :gun: :persevere: :smiling_imp: \39:sweat: :broken_heart: :yellow_heart: :musical_note: :speak_no_evil: \40:wink: :skull: :confounded: :smile: :stuck_out_tongue_winking_eye: \41:angry: :no_good: :muscle: :facepunch: :purple_heart: \42:sparkling_heart: :blue_heart: :grimacing: :sparkles:".split(' ')43 44def top_elements(array, k):45    ind = np.argpartition(array, -k)[-k:]46    return ind[np.argsort(array[ind])][::-1]47 48 49with open("vocabulary.json", 'r') as f:50    vocabulary = json.load(f)51 52st = SentenceTokenizer(vocabulary, 100)53 54emojimodel = torchmoji_emojis("pytorch_model.bin")55 56if USE_GPU:57  emojimodel.to("cuda:0")58 59def deepmojify(sentence, top_n=5, prob_only=False):60    list_emojis = []61    def top_elements(array, k):62        ind = np.argpartition(array, -k)[-k:]63        return ind[np.argsort(array[ind])][::-1]64 65    tokenized, _, _ = st.tokenize_sentences([sentence])66    tokenized = np.array(tokenized).astype(int)  # convert to float first67    if USE_GPU:68        tokenized = torch.tensor(tokenized).cuda()  # then convert to PyTorch tensor69 70    prob = emojimodel.forward(tokenized)[0]71    if not USE_GPU:72        prob = torch.tensor(prob)73    if prob_only:74        return prob75    emoji_ids = top_elements(prob.cpu().numpy(), top_n)76    emojis = map(lambda x: EMOJIS[x], emoji_ids)77    list_emojis.append(emoji.emojize(f"{' '.join(emojis)}", language='alias'))78    # returning the emojis as a list named as list_emojis79    return list_emojis, prob80 81 82model = nn.Sequential(83    nn.Linear(448, 300),  # Increase the number of neurons84    nn.ReLU(),85    nn.BatchNorm1d(300),  # Batch normalization86 87    nn.Linear(300, 300),  # Increase the number of neurons88    nn.ReLU(),89    nn.BatchNorm1d(300),  # Batch normalization90 91    nn.Linear(300, 200),  # Increase the number of neurons92    nn.ReLU(),93    nn.BatchNorm1d(200),  # Batch normalization94 95    nn.Linear(200, 125),  # Increase the number of neurons96    nn.ReLU(),97    nn.BatchNorm1d(125),  # Batch normalization98 99    nn.Linear(125, 2),100    nn.Dropout(0.05)  # Dropout101)102 103model.load_state_dict(torch.load("large.pth", map_location=torch.device('cpu')))104model.eval()105 106@app.route('/infer', methods=['POST'])107def translate():108    data = request.get_json()109 110    TEXT = data['text'].lower()111    probs = deepmojify(TEXT, prob_only=True)112    embedding = sentencemodel.encode(TEXT, convert_to_tensor=True)113    INPUT = torch.cat((probs, embedding))114    output = F.softmax(model(INPUT.view(1, -1)), dim=1)115 116    if output[0][1] > 0.68:117        output = "true"118    else:119        output = "false"120 121    return output122 123 124@app.route('/inferverbose', methods=['POST'])125def translateverbose():126    data = request.get_json()127 128    TEXT = data['text'].lower()129    probs = deepmojify(TEXT, prob_only=True)130    embedding = sentencemodel.encode(TEXT, convert_to_tensor=True)131    INPUT = torch.cat((probs, embedding))132    output = F.softmax(model(INPUT.view(1, -1)), dim=1)133 134    if output[0][1] > 0.4:135        output = "true" + str(output[0][1])136    else:137        output = "false" + str(output[0][0])138 139    return output140 141# Define more routes for other operations like download_model, etc.142if __name__ == "__main__":143   config = Config()144   config.bind = ["0.0.0.0:7860"]  # You can specify the host and port here145   asyncio.run(serve(app, config))