Update app.py

app.py

@@ -1,50 +1,73 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+NSFW Text Descriptor using TF-IDF and Cosine Similarity
+Optimized for modularity, memory efficiency, and Gradio integration.
+"""
+
 import gradio as gr
+import numpy as np
 from datasets import load_dataset
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
-import numpy as np
+from itertools import chain
+from typing import List
+
+
+class NSFWTextMatcher:
+    def __init__(self):
+        self.dataset_sources = [
+            "aifeifei798/DPO_Pairs-Roleplay-NSFW",
+            "Maxx0/sexting-nsfw-adultconten",
+            "QuietImpostor/Claude-3-Opus-Claude-3.5-Sonnnet-9k",
+            "HuggingFaceTB/everyday-conversations-llama3.1-2k",
+            "Chadgpt-fam/sexting_dataset"
+        ]
+        self.all_texts = self._load_all_texts()
+        self.vectorizer = TfidfVectorizer()
+        self.tfidf_matrix = self.vectorizer.fit_transform(self.all_texts)
+
+    def _load_all_texts(self) -> List[str]:
+        texts = []
+        for source in self.dataset_sources:
+            try:
+                dataset = load_dataset(source)
+                for split in dataset:
+                    features = dataset[split].features
+                    if 'text' in features:
+                        texts.extend(dataset[split]['text'])
+                    elif 'content' in features:
+                        texts.extend(dataset[split]['content'])
+            except Exception as e:
+                print(f"[WARN] Failed to load dataset {source}: {e}")
+        return texts
+
+    def find_best_match(self, input_text: str) -> str:
+        input_vector = self.vectorizer.transform([input_text])
+        similarity_scores = cosine_similarity(input_vector, self.tfidf_matrix)
+        best_match_idx = np.argmax(similarity_scores)
+        return self.all_texts[best_match_idx]
+
+
+# Instantiate the matcher once (can be made lazy if needed)
+matcher = NSFWTextMatcher()
+
+
+def generate_text(input_text: str) -> str:
+    if not input_text.strip():
+        return "Please enter a valid input."
+    return matcher.find_best_match(input_text)
+
 
-# Load datasets
-nsfw_datasets = [
-    load_dataset("aifeifei798/DPO_Pairs-Roleplay-NSFW"),
-    load_dataset("Maxx0/sexting-nsfw-adultconten"),
-    load_dataset("QuietImpostor/Claude-3-Opus-Claude-3.5-Sonnnet-9k"),
-    load_dataset("HuggingFaceTB/everyday-conversations-llama3.1-2k"),
-    load_dataset("Chadgpt-fam/sexting_dataset")
-]
-
-# Prepare all texts from datasets
-all_texts = []
-for dataset in nsfw_datasets:
-    for split in dataset.keys():
-        if 'text' in dataset[split].features:
-            all_texts.extend(dataset[split]['text'])
-        elif 'content' in dataset[split].features:
-            all_texts.extend(dataset[split]['content'])
-
-# Create TF-IDF vectorizer
-vectorizer = TfidfVectorizer()
-tfidf_matrix = vectorizer.fit_transform(all_texts)
-
-def find_best_description(input_text):
-    input_vector = vectorizer.transform([input_text])
-    similarities = cosine_similarity(input_vector, tfidf_matrix)
-    most_similar_index = np.argmax(similarities)
-    return all_texts[most_similar_index]
-
-def generate_text(input_text):
-    return find_best_description(input_text)
-
-# Create Gradio interface
+# Gradio Interface
 iface = gr.Interface(
     fn=generate_text,
     inputs=gr.Textbox(label="Enter text to describe"),
-    outputs="text",
+    outputs=gr.Textbox(label="Best Match"),
     title="NSFW Text Descriptor",
-    description="Enter text to find the best description from NSFW datasets.",
-    allow_flagging="never"
+    description="Match your input with the most similar description from NSFW datasets using TF-IDF.",
+    allow_flagging="never",
 )
 
-# Launch the app
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()