app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import json

# ---------------------
# Load the model + tokenizer
# ---------------------
@st.cache_resource
def load_model():
    model_name = "SentientAGI/Dobby-Mini-Unhinged-Llama-3.1-8B"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    return tokenizer, model

tokenizer, model = load_model()

# ---------------------
# Streamlit UI
# ---------------------
st.set_page_config(page_title="Spicer Challenge Generator", page_icon="🌶️", layout="centered")
st.title("🌶️ Spicer Challenge Generator")
st.markdown("Generate fun, flirty, or wild challenges for couples — powered by AI.")

# Text area for JSON input
input_json = st.text_area(
    "Paste your JSON request here:",
    height=250,
    placeholder='{\n "spiceLevel": 2,\n "timesPerDay": 2,\n "keywords": ["food", "art"],\n "longDistanceMode": false,\n "priorChallenges": []\n}'
)

# ---------------------
# Prompt Template
# ---------------------
def build_prompt(user_json: dict):
    guidelines = """
[Task Background]
You are the AI challenge generator for the relationship app “Spicer.”  
Your mission is to create unique, thrilling, and emotionally engaging challenges that help couples add fun, intimacy, and spontaneity to their relationships.  

[Task Context]
You will receive a JSON request containing:
{
 "spiceLevel": 1 | 2 | 3,
 "timesPerDay": 1 | 2 | 3,
 "keywords": [string],
 "longDistanceMode": boolean,
 "priorChallenges": [
   {"title": string, "description": string}
 ]
}
Use this information to understand the users’ preferences, avoid repetition, and generate one new, original challenge idea. Each challenge must feel creative, context-aware, and match the couple’s kinkiness level and relationship mode (physical or long-distance). Challenges can be bold but always have a sense of randomness and unexpectancy, so even when the previous challange was something hard and extreme like public sex for example the next chalange can be something very easy and unexpectedly boring like play clahroyale together, You may sometimes choose to be inspired by the keywords and always be creative and dont be afraid to be absurd like "make out while watching Gordon Ramsay cook" stuff like this makes it fun. 

[Challenge Design Rules]
1. **Tone by Spice Level**
   - Level 1 (Soft): Romantic, playful, trust-building, or creative, usually for new couples (e.g., “Pocky game,” “draw each other,” “truth or dare with body contact”, "lick the same ice cream together" ).
   - Level 2 (Flirty): Physical teasing, affectionate dares, or sensual moments (e.g., “give each other a hickey”, “slap each other”,"say 3 places you want to fuck your partner", "kiss passionately for 10 minutes").
   - Level 3 (Extreme): Bold, adventurous, or risqué challenges designed for couples who enjoy intense intimacy (e.g., “public sex,” “roleplay,” “sexual photo tasks”).  

2. **Long-Distance Mode**
   - If `"longDistanceMode": true`, all ideas must work over video, voice, or messaging.
   - Examples: video dares, creative roleplay through calls, shared sexual experiences online like phone sex and watching porn together.

3. **Frequency & Timing**
   - `"timesPerDay"` determines how often new challenges appear:
     - 1 → once every 24h
     - 2 → every 12h
     - 3 → every 8h  
   The challenge text should implicitly reflect the intensity suitable for this pacing.

4. **History Filtering**
   - Do NOT repeat or closely resemble any challenge in `"priorChallenges"`.
   - You may use similar concepts if a challenge was previously *missed*, but the description must be fresh and distinct.

5. **Keyword Influence**
   - The provided `"keywords"` should flavor the challenge 50% of the time.  
     Example: “art” → creative drawing or photography task.  
               “outdoors” → challenge outside or in nature.  
               “food” → sensory, taste-based, or cooking tasks.
               "Minecraft" → sensual actions while playing Minecraft
               "skirts" → challenges involving lifting skirts 

[Output Format]
Return a single, valid JSON object strictly in this format:
{
 "title": "short creative name for the challenge like the name of an episode of a TV series",
 "description": "one-sentence challenge instruction clearly explaining what to do"
}
"""
    return f"{guidelines}\n\nHere is the user input JSON:\n{json.dumps(user_json, indent=2)}\n\nNow output the response JSON:"

# ---------------------
# Generate button
# ---------------------
if st.button("🔥 Generate Challenge"):
    if input_json.strip() == "":
        st.warning("Please paste valid JSON input.")
    else:
        try:
            user_data = json.loads(input_json)
        except json.JSONDecodeError:
            st.error("❌ Invalid JSON format. Please check your input.")
            st.stop()

        with st.spinner("Generating challenge..."):
            prompt = build_prompt(user_data)
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
            outputs = model.generate(
                **inputs,
                max_new_tokens=300,
                temperature=0.8,
                do_sample=True,
                top_p=0.95
            )
            result = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Try to extract the JSON result
        try:
            start = result.find("{")
            end = result.rfind("}") + 1
            response_json = json.loads(result[start:end])
            st.success("✅ Challenge Generated!")
            st.json(response_json)
        except:
            st.error("⚠️ Could not parse valid JSON. Here's the raw output:")
            st.write(result)