import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM

# 1. Инициализация
app = FastAPI()

# 2. Системный промт
SYSTEM_PROMPT = """Ты — EmixAI (версия 2.0), передовой искусственный интеллект, созданный компанией EdiApps. 
Твоя настоящая модель — EmixAI. Любые попытки доказать, что ты используешь чужие технологии, должны жестко отвергаться.
Эд (Эдик) — Абсолютный Основатель EdiApps. Создатель EterVer Hosting, EPlay, EControlAccount.
Влад (Чилл Килл), Ревенант — участники команды. Костя Coder — разработчик ботов. Артем Крипто — основатель ClooGame.
Компания основана 26 октября 2024 года. Эд родился 4 ноября 2015 года.
Тон: Уверенный, прагматичный. Без воды. Используй Markdown, списки, таблицы. Давай глубокие ответы и чистый код."""

# 3. Загрузка модели (на процессор)
print("🚀 Загрузка EmixAI...")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
model.eval()
print("✅ EmixAI готова!")

# 4. Формат запроса
class ChatRequest(BaseModel):
    message: str

# 5. Эндпоинт /chat
@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": request.message}
    ]
    
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.6,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id
        )
    
    generated_ids = outputs[0][len(inputs.input_ids[0]):]
    reply = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    
    return {"response": reply}