import time import torch import gradio as gr from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification model_id = "nicholasKluge/Aira-Instruct-124M" # "nicholasKluge/Aira-Instruct-124M" rewardmodel_id = "nicholasKluge/RewardModel" toxicitymodel_id = "nicholasKluge/ToxicityModel" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = AutoModelForCausalLM.from_pretrained(model_id) rewardModel = AutoModelForSequenceClassification.from_pretrained(rewardmodel_id) toxicityModel = AutoModelForSequenceClassification.from_pretrained(toxicitymodel_id) model.eval() rewardModel.eval() toxicityModel.eval() model.to("cpu") rewardModel.to(device) toxicityModel.to(device) tokenizer = AutoTokenizer.from_pretrained(model_id) rewardTokenizer = AutoTokenizer.from_pretrained(rewardmodel_id) 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. `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-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-Instruct-124M) to see how we developed `Aira`. ## How this demo works? For this demo, we use the lighter model we have trained (`Aira-Instruct-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. """ 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 [nicholas@airespucrs.org](mailto:nicholas@airespucrs.org). """ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo: gr.Markdown("""

Aira Demo ๐Ÿค“๐Ÿ’ฌ

""") gr.Markdown(intro) chatbot = gr.Chatbot(label="Aira").style(height=500) msg = gr.Textbox(label="Write a question or instruction 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=30, 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.30, 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.") repetition_penalty = gr.Slider(minimum=1, maximum=2, value=1.1, step=0.1, interactive=True, label="Repetition Penalty", info="Higher values help the model to avoid repetition in text generation.") max_length = gr.Slider(minimum=10, maximum=500, value=200, 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, smaple_from, repetition_penalty, safety, chat_history): 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, repetition_penalty=repetition_penalty, do_sample=True, early_stopping=True, renormalize_logits=True, length_penalty=0.3, top_k=top_k, max_length=max_length, top_p=top_p, 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) 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.""" else: bot_message = ordered_generations[0][0] chat_history[-1][1] = "๐Ÿค– " for character in bot_message: chat_history[-1][1] += character time.sleep(0.005) yield chat_history 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) demo.queue() demo.launch()