import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# === Load Personality ===
with open("personality.txt", "r") as f:
    personality = f.read().strip()

# === Model Setup ===
model_name = "UnfilteredAI/DAN-Qwen3-1.7B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,  # Use float32 for CPU compatibility
    device_map="auto"
)

def chat(prompt, history):
    full_prompt = f"{personality}\n\nUser: {prompt}\nDreamer:"
    input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids.to(model.device)

    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.9,
            top_p=0.95
        )

    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return response.split("Dreamer:")[-1].strip()

# === Launch Interface with Public Link ===
gr.ChatInterface(chat).launch()