import gradio as gr
from transformers import GPTJForCausalLM, GPT2Tokenizer

# Load model and tokenizer
model_name = "EleutherAI/gpt-j-6B"
model = GPTJForCausalLM.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)  # GPT-J uses the same tokenizer as GPT-2

# Function to filter explicit content
def filter_explicit(content, filter_on):
    explicit_keywords = ["badword1", "badword2", "badword3"]  # Add more explicit words to filter
    if filter_on:
        for word in explicit_keywords:
            content = content.replace(word, "[CENSORED]")
    return content

def generate_response(prompt, explicit_filter):
    try:
        inputs = tokenizer.encode(prompt, return_tensors="pt")
        outputs = model.generate(
            inputs, 
            max_length=100, 
            num_return_sequences=1,
            temperature=0.7,  # Control the creativity of the response
            top_k=50,         # Limits the sampling pool to top 50 tokens
            top_p=0.9         # Nucleus sampling to avoid repetitive phrases
        )
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        filtered_response = filter_explicit(response, explicit_filter)
        return filtered_response
    except Exception as e:
        return f"Error: {str(e)}"

# Define Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=[gr.Textbox(lines=2, placeholder="Type your message here...", label="Input"), gr.Checkbox(label="Enable Explicit Content Filter")],
    outputs=gr.Textbox(label="Response"),
    title="Chatbot with Explicit Content Filter",
    description="A simple chatbot that allows you to enable or disable explicit content filtering.",
    theme="compact",
    layout="vertical"
)

if __name__ == "__main__":
    iface.launch()