Сбер открывает доступ к нейросетевой модели ruGPT-3.5

# Requirements:
# pip install --user pyzmq

# Port for listening for commands
LOCAL_PORT = 5556
FILE_SETTINGS = 'settings.ini'
FILE_BADWORDS = 'bad_words.txt'
FILE_AUTHORNOTES = 'notes.txt'
# Minimal amount of pure text tokens (without Memory, but with Author Notes)
TEXT_MIN_TOKENS = 10 
# Run on GPU
device='cuda'

import configparser
import zmq
import os
import re
import torch
from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

# Load file if exists, and returns its string content
# If file not exists, returns ""
def load_file(fname):
    if not os.path.isfile(fname):
        return ""
    with open(fname) as f:
        txt = f.read()
    return txt

# Return list of positions of end markers of sentences
def get_ends_of_sentences(text):
    marks = ['\. ','\.\n','\! ','!\n','\? ','\?\n']
    exp = re.compile('|'.join(marks))
    return [m.start() for m in exp.finditer(text)]

# Remove all inbetween < and >, then remove line-endings after #,
# Then remove entire lines beginning with #
def remove_comments(text):
    # remove all inbetween < and > (multiline comment)
    text = re.sub(r'<.+?>', '', text, flags=re.DOTALL)
    # remove any line, where first symbol is #
    text = re.sub(r'^#.+?(\n|$)', '', text, flags=re.MULTILINE)
    # remove all after # (suffix comment)
    text = re.sub(r'#.+?($)', r'\1', text, flags=re.MULTILINE)

    return text

# Load transformer
print("Load model...")
quantize_config = BaseQuantizeConfig(bits=8,
  group_size=128,
  damp_percent=0.01,
  desc_act=False,
  sym=True,
  true_sequential=True)
tokenizer = AutoTokenizer.from_pretrained('model')
model = AutoGPTQForCausalLM.from_quantized('model', device="cuda:0", use_triton=True, quantize_config=quantize_config)

# Open socket
context = zmq.Context()
socket = context.socket(zmq.REP)
socket.bind(f'tcp://*:{LOCAL_PORT}')

config = configparser.ConfigParser()

print("Model loaded")

while True:
    msg = socket.recv_string()

    # Read settings
    config.read_dict({'settings':{ 'top_k':'40',
                                'top_p':'0.98',
                                'repetition_penalty':'1.04',
                                'temperature':'0.6',
                                'num_beams':'1',
                                'num_beam_groups':'1',
                                'max_length':'50',
                                'min_length':'20',
                                'no_repeat_ngram_size':'0',
                                'author_notes_pos':'5',
                                'max_context_tokens':'1024',
                                'show_tokens':'1'},
                     'files':{
                                'memory': 'memory.txt'
                     }})
    # Write default settings if FILE_SETTINGS not exists
    if not os.path.isfile(FILE_SETTINGS):
        with open(FILE_SETTINGS, 'w') as configfile:
            config.write(configfile)
    else:
        config.read(FILE_SETTINGS)

    settings = config['settings']
    min_length = settings.getint('min_length')
    is_lm = True

    # Get message body (remove LM tag for compatibility with T5 plugins)
    if len(msg)>4 and msg[:4]=='<LM>':
        text = msg[4:]
    else:
        text = msg

    # Load Memory text
    file_memory = config['files']['memory']
    memory = remove_comments(load_file(file_memory))

    # Add AuthorNotes to text
    notes_pos = settings.getint('author_notes_pos')
    if notes_pos>0:
        author_notes = load_file(FILE_AUTHORNOTES)
    else:
        author_notes = ""

    if len(author_notes)>0:
        sent_end_list = get_ends_of_sentences(text)
        if len(sent_end_list)<=notes_pos:
            text = author_notes + ' ' + text
        else:
            insert_pos = sent_end_list[-notes_pos] + 2
            text = text[:insert_pos] + author_notes.rstrip() + ' ' + text[insert_pos:]

    # Truncate memory if too big
    max_tokens = settings.getint('max_context_tokens')
    memory_tokenized = tokenizer.encode(memory)
    memory_length = len(memory_tokenized)
    print(f'Tokens in memory file: {memory_length}')
    assert max_tokens>TEXT_MIN_TOKENS,f"max_tokens should be more than {TEXT_MIN_TOKENS}"
    if memory_length > (max_tokens - TEXT_MIN_TOKENS):
        memory_tokenized = memory_tokenized[:max_tokens - TEXT_MIN_TOKENS]

    # Truncate text if too big
    text_tokenized = tokenizer.encode(text)
    text_length = len(text_tokenized)
    text_tokenized = text_tokenized[-(max_tokens-memory_length-1):]

    # Assemble context
    print(tokenizer.decode(memory_tokenized) + tokenizer.decode(text_tokenized))
    lm_text = memory_tokenized + text_tokenized

    # Load stop words list
    stop_list = load_file(FILE_BADWORDS).split('\n')
    stop_list = [s for s in stop_list if len(s)>0]
    if len(stop_list)>0:
        bad_words_ids = tokenizer(stop_list, add_special_tokens=False).input_ids
    else:
        bad_words_ids = None

    input_ids=torch.tensor([lm_text]).to(device)
    num_of_tokens = len(input_ids[0])
    print(f'Num of tokens: {len(input_ids[0])}')

    if settings.getint('num_beam_groups')>1:
        do_sample = False
    else:
        do_sample = True

    attention_mask = torch.ones(input_ids.shape, device=device)
    outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.eos_token_id, bad_words_ids=bad_words_ids, early_stopping=True, do_sample=do_sample, min_length = min_length, max_new_tokens = settings.getint('max_length'), num_beams = settings.getint('num_beams'), num_beam_groups = settings.getint('num_beam_groups'), no_repeat_ngram_size = settings.getint('no_repeat_ngram_size'), repetition_penalty = settings.getfloat('repetition_penalty'), top_p = settings.getfloat('top_p'), top_k = settings.getint('top_k'), temperature = settings.getfloat('temperature') )

    # Print result
    res = tokenizer.decode(outputs[0][num_of_tokens:])
    if settings.getint('show_tokens')>0:
        # Add borders "|" between tokens and show result
        tokens = outputs[0][num_of_tokens:]
        ext_tokens = []
        for t in tokens:
            ext_tokens.append(t)
            ext_tokens.append(97) # token '|' as separator
        print(tokenizer.decode(ext_tokens))
    else:
        print(res)
    print("")

    socket.send_string(res)

# Port for sending commands
LOCAL_PORT = 5556

INPUT_FILE = "input.txt"

import zmq

context = zmq.Context()
socket = context.socket(zmq.REQ)
socket.connect(f'tcp://localhost:{LOCAL_PORT}')

# Send request and get answer
with open(INPUT_FILE) as f:
    req = f.readlines()

socket.send_string("".join(req))

message = socket.recv_string()
print(f'{message}')

# Все строки, содержащие символ решётки (#), являются комментариями и удаляются перед началом обработки файла.
# Структура файла:
# Файл состоит из множества записей. Записи разделены пустой строкой (одной и только одной)
# Каждая запись содержит две строки ровно.
# Первая строка - ключи. Просто фрагменты слов, по которым идёт поиск.
# Ключи разделены запятой, после разделения по запятым пробелы слева и справа каждого ключевика вырезаются strip()
# Вторая строка - собственно содержимое ключа, которое добавляется в контекст, если там встречаются указанные ключи.
# ПРИМЕР:
Демон, демон
[ Демоны - раса могущественных существ. ]

Гоблин, гоблин
[ Гоблины - зеленокожие гуманоиды. ]