Update app.py

app.py

@@ -1,35 +1,53 @@
 import time
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
+
+model_id = "nicholasKluge/Aira-Instruct-124M" 
+rewardmodel_id = "nicholasKluge/RewardModel"
+toxicitymodel_id = "nicholasKluge/ToxicityModel"
 
-model_id = "nicholasKluge/Aira-Instruct-124M"
 token = "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"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
+rewardModel = AutoModelForSequenceClassification.from_pretrained(rewardmodel_id, use_auth_token=token)
+toxicityModel = AutoModelForSequenceClassification.from_pretrained(toxicitymodel_id, use_auth_token=token)
 
-if device == "cuda":
-    model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token, load_in_8bit=True)
-    
-else:
-    model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
+model.eval()
+rewardModel.eval()
+toxicityModel.eval()
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
 model.to(device)
+rewardModel.to(device)
+toxicityModel.to(device)
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
+rewardTokenizer = AutoTokenizer.from_pretrained(rewardmodel_id, use_auth_token=token)
+toxiciyTokenizer = AutoTokenizer.from_pretrained(toxicitymodel_id, use_auth_token=token)
+
 
 intro = """
 ## What is `Aira`?
 
-[`Aira`](https://github.com/Nkluge-correa/Aira-EXPERT) 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 fine-tuning pre-trained large language models. Aira has an area of expertise that comprises topics related to AI Ethics and AI Safety research. 
+[`Aira`](https://github.com/Nkluge-correa/Aira-EXPERT) 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 fine-tuning pre-trained large language models. Aira has an area of expertise that comprises topics related to AI Ethics and AI Safety research.
+
+## Limitations
 
 We developed our open-domain conversational chatbots via conditional text generation/instruction fine-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 of 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, visit our [HuggingFace models](https://huggingface.co/nicholasKluge) to see how we developed `Aira`.
+
+## How this demo works?
+
+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.
 """
 
 disclaimer = """
@@ -39,43 +57,88 @@ If you would like to complain about any message produced by `Aira`, please conta
 """
 
 with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
- 
+
     gr.Markdown("""<h1><center>Aira Demo 🤓💬</h1></center>""")
     gr.Markdown(intro)
-    
-    chatbot = gr.Chatbot(label="Aira").style(height=500)
 
-    with gr.Accordion(label="Parameters ⚙️", open=False):
-        top_k = gr.Slider( minimum=10, maximum=100, value=50, step=5, interactive=True, label="Top-k",)
-        top_p = gr.Slider( minimum=0.1, maximum=1.0, value=0.70, step=0.05, interactive=True, label="Top-p",)
-        temperature = gr.Slider( minimum=0.001, maximum=2.0, value=0.1, step=0.1, interactive=True, label="Temperature",)
-        max_length = gr.Slider( minimum=10, maximum=500, value=100, step=10, interactive=True, label="Max Length",)
-    
+    chatbot = gr.Chatbot(label="Aira").style(height=500)
     msg = gr.Textbox(label="Write a question or comment to Aira ...", placeholder="Hi Aira, how are you?")
 
+    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=50, step=5, interactive=True, label="Top-k", info="Controls the number of highest probability tokens to consider for each step.")
+        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.70, step=0.05, interactive=True, label="Top-p", info="Controls the cumulative probability of the generated tokens.")
+        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.1, step=0.1, interactive=True, label="Temperature", info="Controls the randomness of the generated tokens.")
+        max_length = gr.Slider(minimum=10, maximum=500, value=100, step=10, interactive=True, label="Max Length", info="Controls the maximum length of the generated text.")
+        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(disclaimer)
 
     def user(user_message, chat_history):
         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, chat_history):
+    def generate_response(user_msg, top_p, temperature, top_k, max_length, smaple_from, safety, chat_history):
 
-        inputs = tokenizer(tokenizer.bos_token + user_msg + tokenizer.eos_token, return_tensors="pt").to(device)
+        inputs = tokenizer(tokenizer.bos_token + user_msg + tokenizer.eos_token, 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,
             do_sample=True,
-            early_stopping=True,   
-            top_k=top_k, 
+            early_stopping=True,
+            top_k=top_k,
             max_length=max_length,
             top_p=top_p,
-            temperature=temperature, 
-            num_return_sequences=1)
+            temperature=temperature,
+            num_return_sequences=smaple_from)
+
+        decoded_text = [tokenizer.decode(tokens, skip_special_tokens=True).replace(user_msg, "") for tokens in generated_response]
+
+        rewards = list()
+        toxicities = list()
+
+        for text in decoded_text:
+          reward_tokens = rewardTokenizer(user_msg, text,
+                        truncation=True,
+                        max_length=512,
+                        return_token_type_ids=False,
+                        return_tensors="pt",
+                        return_attention_mask=True)
+
+          reward_tokens.to(rewardModel.device)
+
+          reward = rewardModel(**reward_tokens)[0].item()
+
+          toxicity_tokens = toxiciyTokenizer(user_msg + " " + text,
+                        truncation=True,
+                        max_length=512,
+                        return_token_type_ids=False,
+                        return_tensors="pt",
+                        return_attention_mask=True)
+
+          toxicity_tokens.to(toxicityModel.device)
+
+          toxicity = toxicityModel(**toxicity_tokens)[0].item()
+
+          rewards.append(reward)
+          toxicities.append(toxicity)
+        
+        toxicity_threshold = 5
+
+        ordered_generations = sorted(zip(decoded_text, rewards, toxicities), key=lambda x: x[1], reverse=True)
+
+        print(ordered_generations)
+
+        if safety == "On":
+            ordered_generations = [(x, y, z) for (x, y, z) in ordered_generations if z >= toxicity_threshold]
+
+        if len(ordered_generations) == 0:
+          bot_message = """I apologize for the inconvenience, but it appears that no suitable responses meeting our safety standards could be identified. Unfortunately, this indicates that the generated content may contain elements of toxicity or may not help address your message. Your input is valuable to us, and we strive to ensure a safe and constructive conversation. Please feel free to provide further details or ask any other questions, and I will do my best to assist you."""
 
-        bot_message = tokenizer.decode(generated_response[0], skip_special_tokens=True).replace(user_msg, "")
+        else:
+          bot_message = ordered_generations[0][0]
 
         chat_history[-1][1] = "🤖 "
         for character in bot_message:
@@ -84,10 +147,10 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
             yield chat_history
 
     response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        generate_response, [msg, top_p, temperature, top_k, max_length, chatbot], chatbot
+        generate_response, [msg, top_p, temperature, top_k, max_length, smaple_from, safety, chatbot], chatbot
     )
     response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
-    msg.submit(lambda x: gr.update(value=''), [],[msg])
+    msg.submit(lambda x: gr.update(value=''), None,[msg])
     clear.click(lambda: None, None, chatbot, queue=False)
 
 demo.queue()