Update app.py

app.py

@@ -1,26 +1,33 @@
 import gradio as gr
 import torch, os
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from huggingface_hub import login
 
 def init_model():
     global model, tokenizer
-    # Вставьте сюда ваш токен доступа Hugging Face
     hf_token = os.getenv("HF_TOKEN")
 
-    # Аутентификация с использованием токена
+    if hf_token is None:
+        raise ValueError("Hugging Face token is not set. Please set the HF_TOKEN environment variable.")
+
     login(hf_token, add_to_git_credential=True)
 
-    # Загрузка модели и токенизатора без квантования и без распределения на CPU/диск
     tokenizer = AutoTokenizer.from_pretrained("IlyaGusev/saiga_gemma2_10b", token=hf_token)
+
+    # Настройка квантования
+    quant_config = BitsAndBytesConfig(
+        load_in_8bit=True,  # Использование 8-битного квантования
+        llm_int8_threshold=6.0,  # Порог активации для 8-битных весов
+    )
+
     model = AutoModelForCausalLM.from_pretrained(
         "IlyaGusev/saiga_gemma2_10b",
         token=hf_token,
         torch_dtype=torch.float16,  # Использование float16 для уменьшения потребления памяти
-        device_map=None  # Не использовать автоматическое распределение на CPU/диск
+        device_map="auto",  # Автоматическое распределение модели на GPU
+        quantization_config=quant_config,  # Применение конфигурации квантования
     )
 
-    # Явное перемещение модели на GPU, если доступно
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 
@@ -45,18 +52,16 @@ def generate_response(prompt, max_length=100, temperature=0.7, top_p=0.85, repet
     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="Штраф за повторение")
+        gr.Textbox(lines=2, placeholder="Введите ваш текст здесь..."),
+        gr.Slider(20, 200, step=1, default=100, label="Максимальная длина"),
+        gr.Slider(0.1, 1.0, step=0.1, default=0.7, label="Температура"),
+        gr.Slider(0.1, 1.0, step=0.05, default=0.85, label="Top-p"),
+        gr.Slider(1.0, 2.0, step=0.1, default=1.1, label="Штраф за повторение")
     ],
     outputs="text",
     title="LLM Model Demo",