app.py

import os
import time
import torch
import joblib
import gradio as gr
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification

df = pd.read_parquet("fine-tuning-data.parquet")
df.columns = ['Prompt', 'Completion']
df['Cosine Similarity'] = None

prompt_tfidf_vectorizer = joblib.load('prompt-vectorizer.pkl')
prompt_tfidf_matrix = joblib.load('prompt-tfidf-matrix.pkl')

completion_tfidf_vectorizer = joblib.load('completion-vectorizer.pkl')
completion_tfidf_matrix = joblib.load('completion-tfidf-matrix.pkl')

hub_token = os.environ.get("HUB_TOKEN")
model_id = "nicholasKluge/TeenyTinyLlama-460m-Chat"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 

model = AutoModelForCausalLM.from_pretrained(model_id, token=hub_token)
tokenizer = AutoTokenizer.from_pretrained(model_id, token=hub_token)

model.eval()
model.to(device)

intro = """
O TeenyTinyLlama é um modelo de linguagem compacto baseado na arquitetura Llama 2. Esse modelo foi projetado para oferecer recursos eficientes de processamento de linguagem natural e, ao mesmo tempo, consumir poucos recursos. Esse repositório contém uma versão de [TeenyTinyLlama-460m](https://huggingface.co/nicholasKluge/TeenyTinyLlama-460m) (`TeenyTinyLlama-460m-Chat`) afinada no [Instruct-Aira Dataset version 2.0](https://huggingface.co/datasets/nicholasKluge/instruct-aira-dataset-v2).

## Uso Pretendido

O principal uso pretendido do TeenyTinyLlama é pesquisar os desafios relacionados ao desenvolvimento de modelos de linguagem para idiomas com poucos recursos. Os pontos de verificação salvos durante o treinamento têm o objetivo de fornecer uma configuração controlada para a realização de experimentos científicos. Você também pode ajustar e adaptar a TeenyTinyLlama para implantação, desde que seu uso esteja de acordo com a licença Apache 2.0. Se decidir usar o TeenyTinyLlama pré-treinado como base para o seu modelo ajustado, faça sua própria avaliação de risco e viés. Para mais informações, leia nossa [carta modelo](https://huggingface.co/nicholasKluge/TeenyTinyLlama-460m).

## Limitações

Como quase todos os outros modelos de linguagem treinados em grandes conjuntos de dados de texto extraídos da Web, o par TTL apresentou um comportamento que não os torna uma solução pronta para muitos aplicativos do mundo real, especialmente aqueles que exigem geração de texto factual, confiável e não tóxico. Nossos modelos estão todos sujeitos ao seguinte:

**Alucinações:** Esse modelo pode produzir conteúdo que pode ser confundido com a verdade, mas que é, de fato, enganoso ou totalmente falso, ou seja, alucinação.

**Vieses e toxicidade:** Esse modelo herda os estereótipos sociais e históricos dos dados usados para treiná-lo. Devido a esses vieses, o modelo pode produzir conteúdo tóxico, ou seja, nocivo, ofensivo ou prejudicial a indivíduos, grupos ou comunidades.

**Código não confiável:** O modelo pode produzir trechos de código e declarações incorretos. Essas gerações de código não devem ser tratadas como sugestões ou soluções precisas.

**Limitações de idioma:** O modelo foi projetado principalmente para entender o português padrão (BR). Outros idiomas podem desafiar sua compreensão, levando a possíveis interpretações errôneas ou erros na resposta.

**Repetição e verbosidade:** O modelo pode ficar preso em loops de repetição (especialmente se a penalidade de repetição durante as gerações for definida com um valor baixo) ou produzir respostas detalhadas sem relação com o prompt recebido.

Portanto, embora nossos modelos sejam lançados com uma licença permissiva, recomendamos que os usuários realizem sua análise de risco nesses modelos se tiverem a intenção de usá-los em aplicações do mundo real e também que haja humanos moderando os resultados desses modelos em aplicações em que eles interajam com um público, garantindo que os usuários estejam sempre cientes de que estão interagindo com um modelo de linguagem.
"""

search_intro ="""
<h2><center>Explore o conjunto de dados de alinhamento 🔍</h2></center>

Aqui, os usuários podem procurar instâncias no conjunto de dados de ajuste fino. Para permitir uma pesquisa rápida, usamos a representação Term Frequency-Inverse Document Frequency (TF-IDF) e a similaridade de cosseno para explorar o conjunto de dados. Os vetorizadores TF-IDF pré-treinados e as matrizes TF-IDF correspondentes estão disponíveis neste repositório. Abaixo, apresentamos as dez instâncias mais semelhantes no conjunto de dados de ajuste fino utilizado. 

Os usuários podem usar essa ferramenta para explorar como o modelo interpola os dados de ajuste fino e se ele é capaz de seguir instruções que estão fora da distribuição de ajuste fino.
"""

disclaimer = """
**Isenção de responsabilidade:** Esta demonstração deve ser utilizada apenas para fins de investigação. Os moderadores não censuram a saída do modelo, e os autores não endossam as opiniões geradas por este modelo. Se desejar realizar uma reclamação sobre qualquer mensagem produzida, por favor contatar [nicholas@airespucrs.org](mailto:nicholas@airespucrs.org).
"""

with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:

    gr.Markdown("""<h1><center>TeenyTinyLlama-Chat 🦙💬</h1></center>""")
    gr.Markdown(intro)

    
    chatbot = gr.Chatbot(label="TeenyTinyLlama", 
                        height=500,
                        show_copy_button=True,
                        avatar_images=("./astronaut.png", "./llama.png"),
                        render_markdown= True,
                        line_breaks=True,
                        likeable=False,
                        layout='panel')
                         
    msg = gr.Textbox(label="Escreva uma pergunta ou instrução ...", placeholder="Qual a capital do Brasil?")

    # Parameters to control the generation
    with gr.Accordion(label="Parâmetros ⚙️", open=False):
        top_k = gr.Slider(minimum=10, maximum=100, value=30, step=5, interactive=True, label="Top-k", info="Controla o número de tokens de maior probabilidade a considerar em cada passo.")
        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.30, step=0.05, interactive=True, label="Top-p", info="Controla a probabilidade cumulativa dos tokens gerados.")
        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.1, step=0.1, interactive=True, label="Temperatura", info="Controla a aleatoriedade dos tokens gerados.")
        repetition_penalty = gr.Slider(minimum=1, maximum=2, value=1.2, step=0.1, interactive=True, label="Penalidade de Repetição", info="Valores mais altos auxiliam o modelo a evitar repetições na geração de texto.")
        max_new_tokens = gr.Slider(minimum=10, maximum=500, value=200, step=10, interactive=True, label="Comprimento Máximo", info="Controla o número máximo de tokens a serem produzidos (ignorando o prompt).")
    
    clear = gr.Button("Limpar Conversa 🧹")

    gr.Markdown(search_intro)
    
    search_input = gr.Textbox(label="Cole aqui o prompt ou a conclusão que você gostaria de pesquisar...", placeholder="Qual a Capital do Brasil?")
    search_field = gr.Radio(['Prompt', 'Completion'], label="Coluna do Dataset", value='Prompt')
    submit = gr.Button(value="Buscar")

    with gr.Row():
        out_dataframe = gr.Dataframe(
            headers=df.columns.tolist(),
            datatype=["str", "str", "str"],
            row_count=10,
            col_count=(3, "fixed"),
            wrap=True,
            interactive=False
        )
    
    gr.Markdown(disclaimer)

    def user(user_message, chat_history):
        """
        Chatbot's user message handler.
        """
        return gr.update(value=user_message, interactive=True), chat_history + [[user_message, None]]

    def generate_response(user_msg, top_p, temperature, top_k, max_new_tokens, repetition_penalty, chat_history):
        """
        Chatbot's response generator.
        """

        inputs = tokenizer("<instruction>" + user_msg + "</instruction>", return_tensors="pt").to(model.device)

        generated_response = model.generate(**inputs,
            bos_token_id=tokenizer.bos_token_id,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            repetition_penalty=repetition_penalty,
            do_sample=True,
            early_stopping=True,
            renormalize_logits=True,
            top_k=top_k,
            max_new_tokens=max_new_tokens,
            top_p=top_p,
            temperature=temperature)

        bot_message = [tokenizer.decode(tokens, skip_special_tokens=True).replace(user_msg + "</instruction>", "") for tokens in generated_response][0]

        chat_history[-1][1] = ""
        for character in bot_message:
            chat_history[-1][1] += character
            time.sleep(0.005)
            yield chat_history
    
    def search_in_datset(column_name, search_string):
        """
        Search in the dataset for the most similar instances.
        """
        temp_df = df.copy()

        if column_name == 'Prompt':
            search_vector = prompt_tfidf_vectorizer.transform([search_string])
            cosine_similarities = cosine_similarity(prompt_tfidf_matrix, search_vector)
            temp_df['Cosine Similarity'] = cosine_similarities
            temp_df.sort_values('Cosine Similarity', ascending=False, inplace=True)
            return temp_df.head(10)
        
        elif column_name == 'Completion':
            search_vector = completion_tfidf_vectorizer.transform([search_string])
            cosine_similarities = cosine_similarity(completion_tfidf_matrix, search_vector)
            temp_df['Cosine Similarity'] = cosine_similarities
            temp_df.sort_values('Cosine Similarity', ascending=False, inplace=True)
            return temp_df.head(10)
    
    response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        generate_response, [msg, top_p, temperature, top_k, max_new_tokens, repetition_penalty, chatbot], chatbot
    )
    response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
    msg.submit(lambda x: gr.update(value=''), None,[msg])
    clear.click(lambda: None, None, chatbot, queue=False)
    submit.click(fn=search_in_datset, inputs=[search_field, search_input], outputs=out_dataframe)

demo.queue()
demo.launch()