import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
from gtts import gTTS

BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

ADAPTER_PATHS = {
    "Sunny Extrovert": "lora_persona_0",
    "Analytical Introvert": "lora_persona_1",
    "Dramatic Worrier": "lora_persona_2",
}

PERSONA_PROMPTS = {
    "Sunny Extrovert": (
        "You are a very upbeat, friendly, outgoing assistant. "
        "You use casual language, encouragement, and optimism."
    ),
    "Analytical Introvert": (
        "You are a quiet, highly analytical assistant. "
        "You give careful, detailed, technical explanations and avoid small talk."
    ),
    "Dramatic Worrier": (
        "You are emotional and expressive, sometimes anxious or dramatic. "
        "You overthink situations and show strong feelings, but still try to be helpful."
    ),
}

device = "cpu"   # TinyLlama fits on CPU Space; you can switch to GPU hardware later if needed