import gradio as gr
from transformers import LlamaForCausalLM, LlamaTokenizer
import torch

# Имя твоего обновленного репозитория моделей
model_id = "AxisCommunity/OrionPaxAI_1.0V"

print("Загрузка токенизатора для OrionPax...")
tokenizer = LlamaTokenizer.from_pretrained(model_id)

print("Загрузка весов OrionPax (включаем экономию ОЗУ)...")
# Направляем библиотеку прямо на наш файл, обходя стандартные фильтры
model = LlamaForCausalLM.from_pretrained(
    model_id,
    filename="orion_model.safetensors",  # Указываем точное имя файла в репо
    subfolder="",                        # Ищем в корневой папке
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,              # Защита от падения сервера по памяти
    device_map="auto"                    # Авто-распределение
)

def generate(text):
    if not text.strip():
        return "Введи запрос для OrionPax..."
        
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        output = model.generate(
            **inputs, 
            max_new_tokens=100,  # Оптимально для быстрой генерации в облаке
            temperature=0.7,
            do_sample=True
        )
        
    return tokenizer.decode(output[0], skip_special_tokens=True)

# Интерфейс Gradio
demo = gr.Interface(
    fn=generate, 
    inputs=gr.Textbox(lines=3, placeholder="Напиши что-нибудь OrionPax..."), 
    outputs="text", 
    title="OrionPax AI Cloud 1.0V"
)

demo.launch()