import gradio as gr
from transformers import pipeline
import os

os.environ["HF_HOME"] = "/data/.huggingface"
os.environ["HF_DATASETS_CACHE"] = "/data/.cache/huggingface/datasets"
os.environ["TRANSFORMERS_CACHE"] = "/data/.cache/huggingface/transformers"

print("Loading model...")
pipe = pipeline("text-generation", model="Andrewstivan/AURA", device=-1)
print("Model loaded!")

def chat(message, history):
    result = pipe(
        message,
        max_new_tokens=300,
        temperature=0.8,
        do_sample=True,
        truncation=True,
        max_length=512
    )
    full_text = result[0]["generated_text"]
    return full_text[len(message):].strip()

# Создаём интерфейс
demo = gr.ChatInterface(
    fn=chat,
    title="✨ Aura — поэтичная модель",
    description="Пишите по-русски, модель отвечает на русском"
)

# Запускаем
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)