app.py

from fastapi import FastAPI, Request, Form
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

app = FastAPI()

# Configura las plantillas Jinja2
templates = Jinja2Templates(directory="templates")

# Define el personaje
personaje = "rias"
user="user"
chat={
    personaje: f"hola soy {personaje} no esperaba verte por aqui",
    user:f"hola "
}
# Monta la carpeta 'static' para servir archivos estáticos
app.mount("/static", StaticFiles(directory="static"), name="static")

# Ruta para mostrar los personajes
@app.get("/", response_class=HTMLResponse)
async def read_html(request: Request):
    return templates.TemplateResponse("listapersonajes.html", {"request": request})

# Ruta dinámica para cada personaje
@app.get("/personajes/{personaje}", response_class=HTMLResponse)
async def personaje_detalle(request: Request, personaje: str):
    # El contexto es el nombre de la imagen que se usará
    context = {
        "character_image": f"{personaje}.jpg" , # Asume que el nombre de la imagen es igual al personaje
        "character_name": personaje.capitalize()  # Nombre del personaje con la primera letra en mayúscula

    }
    return templates.TemplateResponse("chat.html", {"request": request, **context})






# Cambia al nuevo modelo
model_name = "allura-org/MoE-Girl_400MA_1BT"

# Inicialización global
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",  # Utiliza automáticamente la GPU si está disponible
    torch_dtype=torch.float16  # Usa FP16 para eficiencia en GPUs
)

@app.post("/personajes/{personaje}/chat", response_class=HTMLResponse)
async def chat_with_character(request: Request, personaje: str, user_input: str):
    # Crear el prompt dinámico con el formato esperado
    prompt = f"""<|im_start|>system
You are {personaje}, a sexy girl who has been dating the user for 2 months.<|im_end|>
<|im_start|>user
{user_input}<|im_end|>
<|im_start|>assistant
"""

    # Tokenizar el prompt
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")

    # Generar la respuesta
    outputs = model.generate(
        **inputs,
        max_new_tokens=500,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        temperature=0.7
    )

    # Decodificar la respuesta
    generated_response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Filtrar la respuesta para extraer solo el texto del asistente
    response_text = generated_response.split("<|im_start|>assistant")[1].strip().split("<|im_end|>")[0].strip()

    # Devolver la respuesta al usuario
    return {"response": response_text}