Update app.py

app.py

@@ -1,73 +1,98 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-NSFW Text Descriptor using TF-IDF and Cosine Similarity
-Optimized for modularity, memory efficiency, and Gradio integration.
-"""
-
+# app.py
 import gradio as gr
-import numpy as np
+from agents.nsfw_agent import NSFWSemanticChatAgent
+from generators.llm_backend import OpenAIBackend, HuggingFaceBackend
+from prompts.nsfw_templates import NSFWPromptTemplate
 from datasets import load_dataset
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.metrics.pairwise import cosine_similarity
-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()
-
+import os
+from typing import List, Tuple
 
-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)
+class NSFWSemanticChatbot:
+    def __init__(self, backend_type: str = "openai"):
+        """Initialize the complete chatbot system"""
+        self.agent = NSFWSemanticChatAgent()
+        self.prompt_template = NSFWPromptTemplate()
+        
+        # Initialize generation backend
+        if backend_type == "openai":
+            api_key = os.getenv("OPENAI_API_KEY")
+            if not api_key:
+                raise ValueError("OpenAI API key required")
+            self.generator = OpenAIBackend(api_key)
+        else:
+            self.generator = HuggingFaceBackend()
+            
+        self._load_dataset()
+        
+    def _load_dataset(self) -> None:
+        """Load and process NSFW dialogue dataset"""
+        try:
+            # Load your NSFW dataset here
+            # Example: dataset = load_dataset("your_nsfw_dataset")
+            # For demonstration, using placeholder data
+            sample_data = [
+                "That's such an interesting perspective...",
+                "I love how you think about these things...",
+                "Tell me more about what you're feeling...",
+                # Add your actual NSFW dialogue samples
+            ]
+            self.agent.build_index(sample_data)
+        except Exception as e:
+            print(f"Dataset loading error: {e}")
+            
+    def generate_response(
+        self, 
+        chat_history: List[Tuple[str, str]], 
+        user_input: str
+    ) -> List[Tuple[str, str]]:
+        """Main response generation pipeline"""
+        
+        if not user_input.strip():
+            return chat_history + [(user_input, "⚠️ Please provide input")]
+            
+        try:
+            # Step 1: Retrieve semantically similar examples
+            retrieved_examples = self.agent.retrieve_context(user_input, k=3)
+            
+            # Step 2: Build contextual prompt
+            prompt = self.prompt_template.build_context_prompt(
+                user_input, chat_history, retrieved_examples
+            )
+            
+            # Step 3: Generate response
+            bot_response = self.generator.generate_response(prompt, max_tokens=150)
+            
+            # Step 4: Update conversation history
+            updated_history = chat_history + [(user_input, bot_response)]
+            
+            return updated_history
+            
+        except Exception as e:
+            error_response = f"System error: {str(e)}"
+            return chat_history + [(user_input, error_response)]
 
+# Initialize chatbot instance
+chatbot = NSFWSemanticChatbot(backend_type="openai")  # or "huggingface"
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=generate_text,
-    inputs=gr.Textbox(label="Enter text to describe"),
-    outputs=gr.Textbox(label="Best Match"),
-    title="NSFW Text Descriptor",
-    description="Match your input with the most similar description from NSFW datasets using TF-IDF.",
-    allow_flagging="never",
-)
+# Create Gradio interface
+def chat_interface(message, history):
+    """Gradio-compatible chat interface"""
+    return chatbot.generate_response(history, message)
 
+# Launch application
 if __name__ == "__main__":
-    iface.launch()
+    demo = gr.ChatInterface(
+        fn=chat_interface,
+        title="🔞 NSFW Semantic Chatbot",
+        description="Advanced conversational AI using semantic embeddings and retrieval-augmented generation",
+        theme="soft",
+        retry_btn="Regenerate Response",
+        undo_btn="Undo Last",
+        clear_btn="Clear Conversation"
+    )
+    
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True
+    )