Changed censor to function instead of class

app.py

@@ -12,7 +12,7 @@ def fn(
     max_length: int = 100,
     temperature: float = 1.5,
     seed: int = 0,
-    censor: bool = True,
+    censor_profanity: bool = True,
 ):
 
     return generate(
@@ -23,7 +23,7 @@ def fn(
             "temperature": temperature,
         },
         seed=seed,
-        censor=censor,
+        censor_profanity=censor_profanity,
     )
 
 

src/generate.py

@@ -1,14 +1,14 @@
 import transformers
 from transformers import set_seed
 
-from src.profanity_filter import ProfanityFilter
+from src.profanity_filter import censor
 
 
 def generate(
     pipeline: transformers.Pipeline,
     pipeline_args: dict,
     seed: int = 0,
-    censor: bool = True,
+    censor_profanity: bool = True,
 ):
 
     set_seed(seed)
@@ -27,7 +27,7 @@ def generate(
     args.update(pipeline_args)
     generated = pipeline(**args)[0]["generated_text"]
 
-    if censor:
-        generated = ProfanityFilter.censor(generated)
+    if censor_profanity:
+        generated = censor(generated)
 
     return generated

src/profanity_filter.py

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