capivarinha_gpu_quantized.py

import numpy as np
import math
import struct
import os
import threading
import torch
from transformers import StoppingCriteria, AutoModelForCausalLM, AutoTokenizer, StoppingCriteriaList, TextIteratorStreamer
from auto_gptq import AutoGPTQForCausalLM

#Vou dar uma acompanhada aqui como tutorial llmpratico 101 pros curiosos do linkedin
#A construção de prompts llama são como você ve abaixo, Intrução, Entrada, Resposta (originalmente em ingles). Originalmente tava uma bagunca de listas virando strings via join, naceu classe nao. 
class Prompt:
    def __init__(self, contexto_maximo=10000):
        
        #Isso aqui da uma traduzida pra token básica, eu to equalizando +- 1 token 3 chars, foi o gpt que falou. pra um contexto de 4000 então, to feliz com 10k chars
        self.contexto_maximo = contexto_maximo
        # Esse template é uma tradução sem vergonha da alpaca, obrigado stanford
        self.template = [
                "Abaixo está uma instrução que descreve uma tarefa",
                "Escreva uma resposta que apropriadamente satisfaça os pedidos",
                "### Instrução:",
                "Seu nome é Capivarinha, você é uma assistente e amiga. Converse naturalmente de forma alegre dando dicas, recomendações e respondendo perguntas, de o máximo de informações que puder se mantendo dentro do tema"
 #               "Se a ultima entrada for 'Cadê minhas capivarinhas' ou algo parecido tem que responder com [gli gli gli gli gli gli gli gli]",
        ]
        #respostas anteriores pra manutenção de contexto, ias gerativas pegam o contexto e geram novas coisas, o chat é só uma implementação
        self.contexto = []
        self.entrada = "### Entrada:"
        self.resposta = "### Resposta:"
        self.tamanhoAtual = len("\n".join(self.template))
        self.resposta_recente = ""

    def adicionar_contexto(self, texto):
        self.contexto.append(texto)
        self.tamanhoAtual = len("\n".join(self.template + self.contexto))
        #se passar do limite, debulha de cima pra baixo até achar a próxima ### Entrada:. Porque? pra a estrutura ficar bunitinha, se nao fica resposta com resposta, entrada sem resposta... pode mudar se for caotico
        if self.tamanhoAtual > self.contexto_maximo:
            self.contexto.pop(0)
            while prmp[0] != self.entrada:
                self.contexto.pop(0)

    def adiciona_entrada(self, entrada):
        #só adiciona 3 linhas na lista que vai ser concatenada com fim de linha pra devolver, basicamente adiciona a linha de entrada, a linha de conteudo entrado e uma linha de resposta pro modelo responder
        self.adicionar_contexto(self.entrada)
        self.adicionar_contexto(entrada)
        self.adicionar_contexto(self.resposta)

    def adiciona_resposta(self, resposta):
        self.resposta_recente = resposta.split(self.resposta)[-1].split(self.entrada)[0].split("### Instrução:")[0].split("\n###")[0].split("###")[0].strip()
        self.adicionar_contexto(self.resposta_recente)

    def limpar_contexto(self):
        #As vezes a bixinha se perde coitada, tem que dar um cntrl alt del 
        self.contexto = []

    def ultima_resposta(self):
        return self.resposta_recente

    #vim do java sim, algum problema?
    def retorna_tamanho_atual(self):
        return self.tamanhoAtual

    def retorna_prompt(self):
        #monta e retorna o prompt da unica forma que realmente importa: string, o resto é pra nao xingarem meus codigo de novo
        #cabei de pensar que podia deixar o contexto ja tokenizadinho né... bobagi, mas assim é mais didático
        return "\n".join(self.template + self.contexto)

def tester(frase, words):
    for word in words:
        if frase.lower().find(word) > -1:
            return True
    return False

#Se prompt nao era classe modelo então não era nem função direito, pipocava em tudo que é pedaço do codigo, mas la vai uns nerd chingar meu codigo então vou embunitar
#gptq só quer dizer que ta quantizado e que roda na gpu... se tiver mais um significado mete um issue ai pra me informar, auto é graça de quem vez a biblioteca
class ModeloAutoGPTQ:
    #quiser por modelo e token em lugares diferentes pode, mas da na mesma
    def __init__(self, criterio_parada, caminho_modelo_tokenizer="./modelo/", t_padding_side="right", t_use_fast=True, nome_modelo="capivarita_gptq_model-4bit-128g",tamanho_contexto=4096,tensores_seguros=True,processador="cuda:0",usar_trion=False):
        #SentencePiece Byte-Pair Encoding, se um dia tiverem sem fazer nada numa terça é uma leitura. 
        self.tokenizer = AutoTokenizer.from_pretrained(caminho_modelo_tokenizer, padding_side=t_padding_side, use_fast=t_use_fast)
        self.modelo = AutoGPTQForCausalLM.from_quantized(caminho_modelo_tokenizer,
                            model_basename=nome_modelo,
                            #isso aqui é de modelo, tem formas como rope (ou superHot) de expandir o contexto, é basicamente o máximo de tokens que cabem. Temos essa visão de sequencialidade dos modelos, mas ela é falsa, vai tudo de uma vez. 
                            max_position_embeddings=tamanho_contexto,
                            #pra ser sincero tava cansado demais pra ler isso aqui, é uma configuração de quantização que protege a convergencia? to alucinando aqui mas pode ter dado sorte de acertar
                            use_safetensors=tensores_seguros,
                            #gpu cuda, esta quantização é só pra gpu, mas deve rodar em quase toda gpu dedicada, só ter uns 4G... nunca testei em 4G, mas imagino. 
                            device=processador,
                            #isso aqui é pros usuarios de maçãzinha, como eu sou usuario de janela né falso. Os linuxeiros ai não sei, acho que da pra usar.
                            use_triton=usar_trion,
                            #isso daqui eu odeio mas to com preguiça, tando none aqui, as infos de quantização tão num documento json junto com o modelo, vão dar uma olhada... ou não tbm
                            quantize_config=None
        )
        self.criterio_parada = criterio_parada
        self.processador = processador

    def generate(self, text_prompt, return_tensors='pt', skip_prompt=True, max_new_tokens=512, repetition_penalty=1.2, temperature=0.9):
        input_ids = self.tokenizer(text_prompt, return_tensors=return_tensors).to(self.processador)
        #mas abandonei o java ta??!!
        self.criterio_parada.comp_inicial = len(input_ids)
        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=skip_prompt)
        generation_kwargs = dict(**input_ids, max_new_tokens=max_new_tokens, streamer=streamer ,repetition_penalty=repetition_penalty, temperature=temperature)
        return threading.Thread(target=self.modelo.generate, kwargs=generation_kwargs), streamer
 


class KeywordsStoppingCriteria(StoppingCriteria):
    def __init__(self, keywords_ids:list):
        self.keywords = keywords_ids
        self.comp_inicial = 0
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        #ve se ja gerou 7 tokens pelo menos novos
        if input_ids[0] > self.comp_inicial + 7:
            #verifica se gerou ###Entrada. "ah ,como vc sabe que são 8 tokens"... hardcoded... olhei e deixei hardcoded.
            for w in self.keywords:
                if tokenizer.decode(input_ids[0][-8:]).find(w) > 0:
                    #para
                    return True
            #numparanaum
        return False

#Inicializa prompt e modelo
prompt = Prompt()
#é só pra parar quando começar a propria capivara fazer a pergunta (entrada) pra ela mesmo... bem TDAH
criterio_parada = KeywordsStoppingCriteria([prompt.entrada])
modelo = ModeloAutoGPTQ(criterio_parada=criterio_parada)


def analizar_entrada(frase):
        #reinicia, dependendo
        if tester(frase,["limpar contexto do monitor", "nova conversa", "reiniciar monitoria"]):
            prompt.limpar_contexto()
            add_message("\n\n[Reiniciado]\n", "italic")
        #se não toma nenhuma ação especial via find, manda pro modelo
        else:
            try:
            #if True:
                prompt.adiciona_entrada(frase)
                modelo_thread, streamer = modelo.generate(prompt.retorna_prompt())
                modelo_thread.start()
                #hacksinho feio pra nao printar a tentativa de entrada, mas to com sono demais pra trocar
                generated_text = ""
                add_message("\nCapivarinha: ", "bold_violet")
                for new_text in streamer:
                    if new_text in ["### Entrada:","### Entrada","### ", "\n###","### ", "\n###","##"]:
                        break
                    generated_text += new_text
                    add_message(new_text)
                add_message("\n")
                prompt.adiciona_resposta(generated_text)
            except:
                add_message("\n\n[Erro de geração]\n", "italic")


def add_message(mes, style=None):
    if len(mes) > 0:
        if mes[-1] == " ":
            print(mes, end="", flush=True)
        else:
            print(mes, end="")
    
def main():
    print("Escreva sua entrada (escreva 'sair' pra sair):\n")
    while True:
        user_input = input("> ")
        if user_input.lower() == "sair":
            break
        analizar_entrada(user_input)

if __name__ == "__main__":
    main()