Update app.py

app.py

@@ -1,26 +1,61 @@
 import time
 import torch
+import joblib
 import gradio as gr
+from datasets import load_dataset
+from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.feature_extraction.text import TfidfVectorizer
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
 
-model_id = "nicholasKluge/Aira-2-124M" # "nicholasKluge/Aira-Instruct-124M"
+
+# download the instruct-aira-dataset
+dataset = load_dataset("nicholasKluge/instruct-aira-dataset", split='english')
+
+# convert the dataset to a pandas dataframe
+df = dataset.to_pandas()
+
+# rename the columns
+df.columns = ['Prompt', 'Completion']
+
+# add a column to store the cosine similarity
+df['Cosine Similarity'] = None
+
+# Load the saved prompt TfidfVectorizer
+prompt_tfidf_vectorizer = joblib.load('prompt_vectorizer.pkl')
+
+# load the prompt tfidf_matrix
+prompt_tfidf_matrix = joblib.load('prompt_tfidf_matrix.pkl')
+
+# Load the saved completion TfidfVectorizer
+completion_tfidf_vectorizer = joblib.load('completion_vectorizer.pkl')
+
+# load the completion tfidf_matrix
+completion_tfidf_matrix = joblib.load('completion_tfidf_matrix.pkl')
+
+# specify the model's ids
+model_id = "nicholasKluge/Aira-2-124M"
 rewardmodel_id = "nicholasKluge/RewardModel"
 toxicitymodel_id = "nicholasKluge/ToxicityModel"
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# specify the device (cuda if available)
+device = torch.device("cpu") # "cuda" if torch.cuda.is_available() else "cpu"
 
+# load the models (chatbot, reward model, toxicity model)
 model = AutoModelForCausalLM.from_pretrained(model_id)
 rewardModel = AutoModelForSequenceClassification.from_pretrained(rewardmodel_id)
 toxicityModel = AutoModelForSequenceClassification.from_pretrained(toxicitymodel_id)
 
+# set the models to evaluation mode
 model.eval()
 rewardModel.eval()
 toxicityModel.eval()
 
-model.to("cpu")
+# set the models to the device
+model.to(device)
 rewardModel.to(device)
 toxicityModel.to(device)
 
+# load the tokenizers
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 rewardTokenizer = AutoTokenizer.from_pretrained(rewardmodel_id)
 toxiciyTokenizer = AutoTokenizer.from_pretrained(toxicitymodel_id)
@@ -28,29 +63,27 @@ toxiciyTokenizer = AutoTokenizer.from_pretrained(toxicitymodel_id)
 
 intro = """
 ## What is `Aira`?
-
 [`Aira`](https://huggingface.co/nicholasKluge/Aira-Instruct-124M) is a `chatbot` designed to simulate the way a human (expert) would behave during a round of questions and answers (Q&A). `Aira` has many iterations, from a closed-domain chatbot based on pre-defined rules to an open-domain chatbot achieved via instruction-tuning.
-
 ## Limitations
-
 We developed our open-domain conversational chatbots via conditional text generation/instruction-tuning. This approach has a lot of limitations. Even though we can make a chatbot that can answer questions about anything, forcing the model to produce good-quality responses is hard. And by good, we mean **factual** and **nontoxic**  text. This leads us to two of the most common problems with generative models used in conversational applications:
-
 🤥 Generative models can perpetuate the generation of pseudo-informative content, that is, false information that may appear truthful.
-
 🤬 In certain types of tasks, generative models can produce harmful and discriminatory content inspired by historical stereotypes.
-
 ## Intended Use
-
 `Aira` is intended only for academic research. For more information, read our [model card](https://huggingface.co/nicholasKluge/Aira-2-124M) to see how we developed `Aira`.
-
 ## How this demo works?
-
 For this demo, we use the lighter model we have trained (`Aira-2-124M`). This demo employs a [`reward model`](https://huggingface.co/nicholasKluge/RewardModel) and a [`toxicity model`](https://huggingface.co/nicholasKluge/ToxicityModel) to evaluate the score of each candidate's response, considering its alignment with the user's message and its level of toxicity. The generation function arranges the candidate responses in order of their reward scores and eliminates any responses deemed toxic or harmful. Subsequently, the generation function returns the candidate response with the highest score that surpasses the safety threshold, or a default message if no safe candidates are identified.
 """
 
+search_intro ="""
+<h2><center>Explore Aira's Dataset 🔍</h2></center>
+
+Here, users can look for instances in Aira's fine-tuning dataset where a given prompt or completion resembles an instruction in its dataset. To enable a fast search, we use the Term Frequency-Inverse Document Frequency (TF-IDF) representation and cosine similarity to explore the dataset. The pre-trained TF-IDF vectorizers and corresponding TF-IDF matrices are available in this repository. Below, we present the top five most similar instances in Aira's dataset for every search query.
+
+Users can use this to explore how the model interpolates on the fine-tuning data and if it is capable of following instructions that are out of the fine-tuning distribution.
+"""
+
 disclaimer = """
 **Disclaimer:** You should use this demo for research purposes only. Moderators do not censor the model output, and the authors do not endorse the opinions generated by this model.
-
 If you would like to complain about any message produced by `Aira`, please contact [[email protected]](mailto:[email protected]).
 """
 
@@ -62,6 +95,7 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
     chatbot = gr.Chatbot(label="Aira").style(height=500)
     msg = gr.Textbox(label="Write a question or instruction to Aira ...", placeholder="What is the capital of Brazil?")
 
+    # Parameters to control the generation
     with gr.Accordion(label="Parameters ⚙️", open=True):
         safety = gr.Radio(["On", "Off"], label="Guard Rail 🛡️", value="On", info="Helps prevent the model from generating toxic/harmful content.")
         top_k = gr.Slider(minimum=10, maximum=100, value=30, step=5, interactive=True, label="Top-k", info="Controls the number of highest probability tokens to consider for each step.")
@@ -72,12 +106,36 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
         smaple_from = gr.Slider(minimum=2, maximum=10, value=2, step=1, interactive=True, label="Sample From", info="Controls the number of generations that the reward model will sample from.")
 
     clear = gr.Button("Clear Conversation 🧹")
+
+    gr.Markdown(search_intro)
+    
+    search_input = gr.Textbox(label="Paste here the prompt or completion you would like to search ...", placeholder="What is your name?")
+    search_field = gr.Radio(['Prompt', 'Completion'], label="Dataset Column", value='Prompt')
+    submit = gr.Button(value="Search")
+
+    with gr.Row():
+        out_dataframe = gr.Dataframe(
+            headers=df.columns.tolist(),
+            datatype=["str", "str", "str"],
+            max_rows=5,
+            col_count=(3, "fixed"),
+            overflow_row_behaviour='paginate',
+            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_length, smaple_from, repetition_penalty, safety, chat_history):
+        """
+        Chatbot's response generator.
+        """
 
         inputs = tokenizer(tokenizer.bos_token + user_msg + tokenizer.sep_token, return_tensors="pt").to(model.device)
 
@@ -146,12 +204,36 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
             time.sleep(0.005)
             yield chat_history
 
+    def search_in_datset(column_name, search_string):
+        """
+        Search in the dataset for the most similar instances.
+        """
+        print(column_name, search_string)
+        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()
+        
+        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()
+
+    
+    
     response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         generate_response, [msg, top_p, temperature, top_k, max_length, smaple_from, repetition_penalty, safety, 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()
+demo.launch()