inference.py

import torch
import torch.nn as nn

itw = {
    0: '\t',
    1: ' ',
    2: '!',
    3: '"',
    4: '#',
    5: '$',
    6: '%',
    7: '&',
    8: '(',
    9: ')',
    10: '*',
    11: '+',
    12: ',',
    13: '-',
    14: '.',
    15: '/',
    16: '0',
    17: '1',
    18: '2',
    19: '3',
    20: '4',
    21: '5',
    22: '6',
    23: '7',
    24: '8',
    25: '9',
    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: '\u200b'
}

wti = {
    '\t': 0,
    ' ': 1,
    '!': 2,
    '"': 3,
    '#': 4,
    '$': 5,
    '%': 6,
    '&': 7,
    '(': 8,
    ')': 9,
    '*': 10,
    '+': 11,
    ',': 12,
    '-': 13,
    '.': 14,
    '/': 15,
    '0': 16,
    '1': 17,
    '2': 18,
    '3': 19,
    '4': 20,
    '5': 21,
    '6': 22,
    '7': 23,
    '8': 24,
    '9': 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,
    '\u200b': 115,
}

class AI(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(len(wti), 128)
        self.lstm = nn.LSTM(128, 128, batch_first=True)
        self.attention = nn.MultiheadAttention(embed_dim=128, num_heads=4, batch_first=True)
        self.layernorm_after = nn.LayerNorm(128)
        self.linear = nn.Linear(128, len(itw))
    def forward(self, x):
        x = self.embed(x)
        x_origin = x
        x, _ = self.lstm(x)
        x_origin = x
        x, _ = self.attention(x, x, x)
        x+= x_origin
        x = self.layernorm_after(x)
        x = self.linear(x)
        x = x[:, -1, :]
        return x
model = AI()
model.to('cuda')
model.load_state_dict(torch.load('patch'))
for _ in range(100):
    a = input()
    a = list(a)
    ss = []
    for i in a:
        try:
            ss.append(wti[i])
        except:
            ss.append(wti[' '])
    for i in range(100):
        sss = torch.tensor(ss).unsqueeze(0)
        otv = model(sss)
        otv = torch.argmax(otv, dim=1).item()
        ss.append(otv)
    txtx = ''
    for i in ss:
        try:
            txtx += itw[i]
        except:
            pass
    print(txtx)
    ss = []