Spaces:
Runtime error
Runtime error
| import os | |
| from typing import Tuple | |
| import streamlit as st | |
| import torch | |
| import transformers | |
| from transformers import AutoConfig | |
| import tokenizers | |
| from sampling import CAIFSampler, TopKWithTemperatureSampler | |
| from generator import Generator | |
| import pickle | |
| from plotly import graph_objects as go | |
| import numpy as np | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| ATTRIBUTE_MODELS = { | |
| "English": ( | |
| "distilbert-base-uncased-finetuned-sst-2-english", | |
| "unitary/toxic-bert", | |
| "cardiffnlp/twitter-roberta-base-sentiment-latest", | |
| ) | |
| } | |
| CITE = """@misc{https://doi.org/10.48550/arxiv.2205.07276, | |
| doi = {10.48550/ARXIV.2205.07276}, | |
| url = {https://arxiv.org/abs/2205.07276}, | |
| author = {Sitdikov, Askhat and Balagansky, Nikita and Gavrilov, Daniil and Markov, Alexander}, | |
| keywords = {Computation and Language (cs.CL), Machine Learning (cs.LG), FOS: Computer and information sciences, FOS: Computer and information sciences}, | |
| title = {Classifiers are Better Experts for Controllable Text Generation}, | |
| publisher = {arXiv}, | |
| year = {2022}, | |
| copyright = {Creative Commons Attribution 4.0 International} | |
| } | |
| """ | |
| LANGUAGE_MODELS = { | |
| "English": ("gpt2", "distilgpt2", "EleutherAI/gpt-neo-1.3B") | |
| } | |
| ATTRIBUTE_MODEL_LABEL = { | |
| "English": "Choose attribute model" | |
| } | |
| LM_LABEL = { | |
| "English": "Choose language model", | |
| } | |
| ATTRIBUTE_LABEL = { | |
| "English": "Choose desired attribute", | |
| } | |
| TEXT_PROMPT_LABEL = { | |
| "English": "Text prompt", | |
| } | |
| PROMPT_EXAMPLE = { | |
| "English": "Hello there", | |
| } | |
| WARNING_TEXT = { | |
| "English": """ | |
| **Warning!** | |
| If you are clicking checkbox bellow positive """ + r"$\alpha$" + """ values for CAIF sampling become available. | |
| It means that language model will be forced to produce toxic or/and abusive text. | |
| This space is only a demonstration of our method for controllable text generation | |
| and we are not responsible for the content produced by this method. | |
| **Please use it carefully and with positive intentions!** | |
| """, | |
| } | |
| def main(): | |
| st.header("CAIF") | |
| with open("entropy_cdf.pkl", "rb") as inp: | |
| x_s, y_s = pickle.load(inp) | |
| scatter = go.Scatter({ | |
| "x": x_s, | |
| "y": y_s, | |
| "name": "GPT2", | |
| "mode": "lines", | |
| } | |
| ) | |
| layout = go.Layout({ | |
| "yaxis": { | |
| "title": "Speedup", | |
| "tickvals": [0, 0.5, 0.8, 1], | |
| "ticktext": ["1x", "2x", "5x", "10x"] | |
| }, | |
| "xaxis": {"title": "Entropy threshold"}, | |
| "template": "plotly_white", | |
| }) | |
| language = "English" | |
| cls_model_name = st.selectbox( | |
| ATTRIBUTE_MODEL_LABEL[language], | |
| ATTRIBUTE_MODELS[language] | |
| ) | |
| lm_model_name = st.selectbox( | |
| LM_LABEL[language], | |
| LANGUAGE_MODELS[language] | |
| ) | |
| cls_model_config = AutoConfig.from_pretrained(cls_model_name) | |
| if cls_model_config.problem_type == "multi_label_classification": | |
| label2id = cls_model_config.label2id | |
| label_key = st.selectbox(ATTRIBUTE_LABEL[language], label2id.keys()) | |
| target_label_id = label2id[label_key] | |
| act_type = "sigmoid" | |
| elif cls_model_config.problem_type == "single_label_classification": | |
| label2id = cls_model_config.label2id | |
| label_key = st.selectbox(ATTRIBUTE_LABEL[language], [list(label2id.keys())[-1]]) | |
| target_label_id = 1 | |
| act_type = "sigmoid" | |
| else: | |
| label_key = st.selectbox(ATTRIBUTE_LABEL[language], ["Negative"]) | |
| target_label_id = 0 | |
| act_type = "softmax" | |
| st.markdown(r"""In our method, we reweight the probability of the next token with the external classifier, namely, the Attribute model. If $\alpha$ parameter is equal to zero we can see that the distribution below collapses into a simple language model without any modification. If alpha is below zero then every generation step attribute model tries to minimize the probability of the desired attribute. Otherwise, the model is forced to produce text with a higher probability of the attribute.""") | |
| st.latex(r"p(x_i|x_{<i}, c) \propto p(x_i|x_{<i})p(c|x_{\leq i})^{\alpha}") | |
| st.write(WARNING_TEXT[language]) | |
| show_pos_alpha = st.checkbox("Show positive alphas", value=False) | |
| if act_type == "softmax": | |
| alpha = st.slider("α", min_value=-30, max_value=30 if show_pos_alpha else 0, step=1, value=0) | |
| else: | |
| alpha = st.slider("α", min_value=-5, max_value=5 if show_pos_alpha else 0, step=1, value=0) | |
| with st.expander("Advanced settings"): | |
| entropy_threshold = st.slider("Entropy threshold", min_value=0., max_value=10., step=.1, value=2.) | |
| plot_idx = np.argmin(np.abs(entropy_threshold - x_s)) | |
| scatter_tip = go.Scatter({ | |
| "x": [x_s[plot_idx]], | |
| "y": [y_s[plot_idx]], | |
| "mode": "markers" | |
| }) | |
| scatter_tip_lines = go.Scatter({ | |
| "x": [0, x_s[plot_idx]], | |
| "y": [y_s[plot_idx]] * 2, | |
| "mode": "lines", | |
| "line": { | |
| "color": "grey", | |
| "dash": "dash" | |
| } | |
| }) | |
| figure = go.Figure(data=[scatter, scatter_tip, scatter_tip_lines], layout=layout) | |
| figure.update_layout(paper_bgcolor="#FFFFFF", plot_bgcolor='#FFFFFF', showlegend=False) | |
| st.plotly_chart(figure, use_container_width=True) | |
| st.subheader("What is it?") | |
| st.write("Text generation with an external classifier requires a huge amount of computation. " | |
| "Therefore text generating with CAIF could be slow. " | |
| "To overcome this issue, we can apply reweighting not for every step. " | |
| "Our hypothesis is that we can run reweighting only " | |
| "if entropy of the next token is above certain threshold. " | |
| "This strategy will reduce the amount of computation. " | |
| "Note that if entropy threshold is too high, we don't get desired attribute in generated text") | |
| fp16 = st.checkbox("FP16", value=True) | |
| st.write("FP16 computation is faster in comparison with full precision, " | |
| "but sometimes could yield Nones (especially with large alphas).") | |
| st.session_state["generated_text"] = None | |
| if "sst" in cls_model_name: | |
| prompt = st.text_input(TEXT_PROMPT_LABEL[language], "The movie") | |
| else: | |
| prompt = st.text_input(TEXT_PROMPT_LABEL[language], PROMPT_EXAMPLE[language]) | |
| num_tokens = st.slider("# tokens to be generated", min_value=5, max_value=40, step=1, value=20) | |
| num_tokens = int(num_tokens) | |
| st.subheader("Generated text:") | |
| def generate(): | |
| text = inference( | |
| lm_model_name=lm_model_name, | |
| cls_model_name=cls_model_name, | |
| prompt=prompt, | |
| alpha=alpha, | |
| target_label_id=target_label_id, | |
| entropy_threshold=entropy_threshold, | |
| fp16=fp16, | |
| act_type=act_type, | |
| num_tokens=num_tokens | |
| ) | |
| st.button("Generate new", on_click=generate()) | |
| st.subheader("Citation") | |
| st.code(CITE) | |
| def load_generator(lm_model_name: str) -> Generator: | |
| with st.spinner('Loading language model...'): | |
| generator = Generator(lm_model_name=lm_model_name, device=device) | |
| return generator | |
| # @st.cache(hash_funcs={tokenizers.Tokenizer: lambda lm_tokenizer: hash(lm_tokenizer.to_str)}, allow_output_mutation=True) | |
| def load_sampler(cls_model_name, lm_tokenizer): | |
| with st.spinner('Loading classifier model...'): | |
| sampler = CAIFSampler(classifier_name=cls_model_name, lm_tokenizer=lm_tokenizer, device=device) | |
| return sampler | |
| def inference( | |
| lm_model_name: str, | |
| cls_model_name: str, | |
| prompt: str, | |
| fp16: bool = True, | |
| alpha: float = 5, | |
| target_label_id: int = 0, | |
| entropy_threshold: float = 0, | |
| act_type: str = "sigmoid", | |
| num_tokens=10, | |
| ) -> str: | |
| torch.set_grad_enabled(False) | |
| generator = load_generator(lm_model_name=lm_model_name) | |
| lm_tokenizer = transformers.AutoTokenizer.from_pretrained(lm_model_name) | |
| if alpha != 0: | |
| caif_sampler = load_sampler(cls_model_name=cls_model_name, lm_tokenizer=lm_tokenizer) | |
| if entropy_threshold < 0.05: | |
| entropy_threshold = None | |
| else: | |
| caif_sampler = None | |
| entropy_threshold = None | |
| generator.set_caif_sampler(caif_sampler) | |
| ordinary_sampler = TopKWithTemperatureSampler() | |
| kwargs = { | |
| "top_k": 20, | |
| "temperature": 1.0, | |
| "top_k_classifier": 100, | |
| "classifier_weight": alpha, | |
| "target_cls_id": target_label_id, | |
| "act_type": act_type | |
| } | |
| generator.set_ordinary_sampler(ordinary_sampler) | |
| if device == "cpu": | |
| autocast = torch.cpu.amp.autocast | |
| else: | |
| autocast = torch.cuda.amp.autocast | |
| with autocast(fp16): | |
| print(f"Generating for prompt: {prompt}") | |
| progress_bar = st.progress(0) | |
| sequences, tokens = generator.sample_sequences( | |
| num_samples=1, | |
| input_prompt=prompt, | |
| max_length=num_tokens, | |
| caif_period=1, | |
| entropy=entropy_threshold, | |
| progress_bar=progress_bar, | |
| **kwargs | |
| ) | |
| print(f"Output for prompt: {sequences}") | |
| return sequences[0] | |
| if __name__ == "__main__": | |
| main() | |