app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import login

def init_model():
    global model, tokenizer
    # Вставьте сюда ваш токен доступа Hugging Face
    hf_token = os.getenv("HF_TOKEN")

    # Аутентификация с использованием токена
    login(hf_token, add_to_git_credential=True)

    # Загрузка модели и токенизатора без квантования и без распределения на CPU/диск
    tokenizer = AutoTokenizer.from_pretrained("IlyaGusev/saiga_gemma2_10b", token=hf_token)
    model = AutoModelForCausalLM.from_pretrained(
        "IlyaGusev/saiga_gemma2_10b",
        token=hf_token,
        torch_dtype=torch.float16,  # Использование float16 для уменьшения потребления памяти
        device_map=None  # Не использовать автоматическое распределение на CPU/диск
    )

    # Явное перемещение модели на GPU, если доступно
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

def generate_response(prompt, max_length=100, temperature=0.7, top_p=0.85, repetition_penalty=1.1):
    try:
        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
        attention_mask = torch.ones_like(input_ids).to(model.device)

        output = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_length=max_length,
            temperature=temperature,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            do_sample=True,
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id
        )
        response_text = tokenizer.decode(output[0], skip_special_tokens=True)
        return response_text
    except Exception as e:
        return f"Извините, произошла ошибка при генерации ответа: {str(e)}"

# Инициализация модели и токенизатора
init_model()

# Создание интерфейса Gradio
iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.inputs.Textbox(lines=2, placeholder="Введите ваш текст здесь..."),
        gr.inputs.Slider(20, 200, step=1, default=100, label="Максимальная длина"),
        gr.inputs.Slider(0.1, 1.0, step=0.1, default=0.7, label="Температура"),
        gr.inputs.Slider(0.1, 1.0, step=0.05, default=0.85, label="Top-p"),
        gr.inputs.Slider(1.0, 2.0, step=0.1, default=1.1, label="Штраф за повторение")
    ],
    outputs="text",
    title="LLM Model Demo",
    description="Введите текстовый запрос, чтобы сгенерировать ответ с помощью LLM модели."
)

if __name__ == "__main__":
    iface.launch()