Upload 6 files

app.py

@@ -3,9 +3,21 @@ import time
 import torch
 import joblib
 import gradio as gr
-from datasets import load_dataset
+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-160m-Chat"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
@@ -40,6 +52,14 @@ Desenvolvemos este modelo de conversação através de ajuste fino por instruç
 TeenyTinyLlama destina-se apenas à investigação academica. Para mais informações, leia nossa [carta modelo](https://huggingface.co/nicholasKluge/TeenyTinyLlama-160m).
 """
 
+search_intro ="""
+<h2><center>Explore o conjunto de dados da Aira 🔍</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.
 
@@ -72,6 +92,22 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
         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)
 
@@ -107,7 +143,27 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
         for character in bot_message:
             chat_history[-1][1] += character
             time.sleep(0.005)
-            yield chat_history  
+            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
@@ -115,6 +171,7 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
     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()

completion-tfidf-matrix.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:713c646967c5dff5f0c66c526553b618f7d7b9ac2d8d580ec1615d2d28211f30
+size 116332899

completion-vectorizer.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64bbd49731efda1677dca239d838058ef221de4770ed5663e451af1347c30637
+size 4036115

fine-tuning-data.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0eb2d74b25cf773e7a2edbcdc3d05818e80eaa494d124e4deb5820158958d7c
+size 89839763

prompt-tfidf-matrix.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b1c4c26ba191e97766b0c1d213277cd212cd574adcdd94cc09494a70d74b5f2
+size 37979435

prompt-vectorizer.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29a72f7e1c286ffe7c87c5384484f5471baec4b3a86238b458fdbcfb52d01a38
+size 3324940