Added dedicated generation function and profanity filter

app.py

@@ -1,72 +1,46 @@
 import gradio as gr
-from transformers import set_seed
 
+from src.generate import generate
 from src.utils import load_pipelines_from_config
 
 pipelines = load_pipelines_from_config(config_path="model_config.yaml")
 
-def generate(
-    text: str,
+
+def fn(
+    text_inputs: str,
     model: str,
     max_length: int = 100,
     temperature: float = 1.5,
-    seed: int = 0
+    seed: int = 0,
+    censor: bool = True,
 ):
-    num_beams: int = 5
-    num_return_sequences: int = 1
-    no_repeat_ngram_size: int = 3
-    early_stopping: bool = True
-    skip_special_tokens: bool = True
-
-    set_seed(seed)
 
-    generated = pipelines[model](
-        text_inputs=text,
-        max_length=max_length,
-        num_return_sequences=num_return_sequences,
-        num_beams=num_beams,
-        no_repeat_ngram_size=no_repeat_ngram_size,
-        early_stopping=early_stopping,
-        skip_special_tokens=skip_special_tokens,
-        temperature=temperature
+    return generate(
+        pipeline=pipelines[model],
+        pipeline_args={
+            "text_inputs": text_inputs,
+            "max_length": max_length,
+            "temperature": temperature,
+        },
+        seed=seed,
+        censor=censor,
     )
 
-    return generated[0]["generated_text"]
-
 
 iface = gr.Interface(
-    fn=generate,
+    fn=fn,
     inputs=[
-        gr.Textbox(
-            value="[Verse]",
-            placeholder="Input text...",
-            label="Input Text"
-        ),
+        gr.Textbox(value="[Verse]", placeholder="Input text...", label="Input Text"),
         gr.Dropdown(
             choices=list(pipelines.keys()),
             value=list(pipelines.keys())[0],
-            label="Model"
-        ),
-        gr.Slider(
-            minimum=50,
-            maximum=1000,
-            value=100,
-            step=10,
-            label="Max Length"
-        ),
-        gr.Slider(
-            minimum=0.4,
-            maximum=1.9,
-            value=1.5,
-            step=0.1,
-            label="Temperature"
-        ),
-        gr.Number(
-            value=0,
-            precision=0,
-            label="Seed"
+            label="Model",
         ),
+        gr.Slider(minimum=50, maximum=1000, value=100, step=10, label="Max Length"),
+        gr.Slider(minimum=0.4, maximum=1.9, value=1.5, step=0.1, label="Temperature"),
+        gr.Number(value=0, precision=0, label="Seed"),
+        gr.CheckBox(value=True, label="Censor Profanity"),
     ],
-    outputs="text"
-    )
+    outputs="text",
+)
 iface.launch()

src/generate.py

@@ -0,0 +1,33 @@
+import transformers
+from transformers import set_seed
+
+from src.profanity_filter import ProfanityFilter
+
+
+def generate(
+    pipeline: transformers.Pipeline,
+    pipeline_args: dict,
+    seed: int = 0,
+    censor: bool = True,
+):
+
+    set_seed(seed)
+
+    default_pipline_args = {
+        "text_inputs": None,
+        "max_length": 100,
+        "num_return_sequences": 1,
+        "num_beams": 5,
+        "no_repeat_ngram_size": 3,
+        "early_stopping": True,
+        "skip_special_tokens": True,
+        "temperature": 1.5,
+    }
+
+    args = default_pipline_args.update(pipeline_args)
+    generated = pipeline(**args)[0]["generated_text"]
+
+    if censor:
+        generated = ProfanityFilter.censor(generated)
+
+    return generated

src/profanity_filter.py

@@ -0,0 +1,49 @@
+import string
+
+import requests
+
+
+class ProfanityFilter:
+    def __init__(self):
+        BANNED_LIST_URL = "https://raw.githubusercontent.com/snguyenthanh/better_profanity/master/better_profanity/profanity_wordlist.txt"
+        self.banned_list = requests.get(BANNED_LIST_URL).text.split("\n")
+
+    def censor(self, text="", censor_char="*", keep_first_letter=True):
+
+        # Split sentences by newline
+        sentence_list = text.split("\n")
+        for s, sentence in enumerate(sentence_list):
+
+            # Split words in sentence by space
+            word_list = sentence.split()
+            for w, word in enumerate(word_list):
+
+                # Process word to match banned list
+                processed_word = word.translate(
+                    str.maketrans("", "", string.punctuation)
+                ).lower()
+
+                # Replace if word is profane
+                if processed_word in self.banned_list:
+                    censored_word = censor_char * len(word)
+
+                    # Keep first letter of word for context if desired
+                    if keep_first_letter:
+                        censored_word = word[0] + censored_word[1:]
+
+                    # Replcate punctuation
+                    censored_word_punc = ""
+                    for c, char in enumerate(word):
+                        if char in string.punctuation:
+                            censored_word_punc += word[c]
+                        else:
+                            censored_word_punc += censored_word[c]
+
+                    # Update word list
+                    word_list[w] = censored_word_punc
+
+                # Update sentence list
+                sentence_list[s] = word_list
+
+        # Join everything back together
+        return "\n".join([" ".join(word_list) for word_list in sentence_list])

src/utils.py

@@ -1,16 +1,17 @@
 import yaml
 from transformers import pipeline
 
+
 def load_pipelines_from_config(config_path: str):
     with open(config_path, "r") as f:
         model_config = yaml.safe_load(f.read())
-        
+
     models = {}
     for model, config in model_config.items():
         models[model] = pipeline(
             task=config["task"],
             model=config["model_name"],
-            revision=config["hf_commit_hash"]
+            revision=config["hf_commit_hash"],
         )
-        
-    return models
+
+    return models