PitterTMYT commited on
Commit
2965f99
·
verified ·
1 Parent(s): 6150bd3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -13
app.py CHANGED
@@ -1,26 +1,33 @@
1
  import gradio as gr
2
  import torch, os
3
- from transformers import AutoTokenizer, AutoModelForCausalLM
4
  from huggingface_hub import login
5
 
6
  def init_model():
7
  global model, tokenizer
8
- # Вставьте сюда ваш токен доступа Hugging Face
9
  hf_token = os.getenv("HF_TOKEN")
10
 
11
- # Аутентификация с использованием токена
 
 
12
  login(hf_token, add_to_git_credential=True)
13
 
14
- # Загрузка модели и токенизатора без квантования и без распределения на CPU/диск
15
  tokenizer = AutoTokenizer.from_pretrained("IlyaGusev/saiga_gemma2_10b", token=hf_token)
 
 
 
 
 
 
 
16
  model = AutoModelForCausalLM.from_pretrained(
17
  "IlyaGusev/saiga_gemma2_10b",
18
  token=hf_token,
19
  torch_dtype=torch.float16, # Использование float16 для уменьшения потребления памяти
20
- device_map=None # Не использовать автоматическое распределение на CPU/диск
 
21
  )
22
 
23
- # Явное перемещение модели на GPU, если доступно
24
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
25
  model.to(device)
26
 
@@ -45,18 +52,16 @@ def generate_response(prompt, max_length=100, temperature=0.7, top_p=0.85, repet
45
  except Exception as e:
46
  return f"Извините, произошла ошибка при генерации ответа: {str(e)}"
47
 
48
- # Инициализация модели и токенизатора
49
  init_model()
50
 
51
- # Создание интерфейса Gradio
52
  iface = gr.Interface(
53
  fn=generate_response,
54
  inputs=[
55
- gr.inputs.Textbox(lines=2, placeholder="Введите ваш текст здесь..."),
56
- gr.inputs.Slider(20, 200, step=1, default=100, label="Максимальная длина"),
57
- gr.inputs.Slider(0.1, 1.0, step=0.1, default=0.7, label="Температура"),
58
- gr.inputs.Slider(0.1, 1.0, step=0.05, default=0.85, label="Top-p"),
59
- gr.inputs.Slider(1.0, 2.0, step=0.1, default=1.1, label="Штраф за повторение")
60
  ],
61
  outputs="text",
62
  title="LLM Model Demo",
 
1
  import gradio as gr
2
  import torch, os
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
4
  from huggingface_hub import login
5
 
6
  def init_model():
7
  global model, tokenizer
 
8
  hf_token = os.getenv("HF_TOKEN")
9
 
10
+ if hf_token is None:
11
+ raise ValueError("Hugging Face token is not set. Please set the HF_TOKEN environment variable.")
12
+
13
  login(hf_token, add_to_git_credential=True)
14
 
 
15
  tokenizer = AutoTokenizer.from_pretrained("IlyaGusev/saiga_gemma2_10b", token=hf_token)
16
+
17
+ # Настройка квантования
18
+ quant_config = BitsAndBytesConfig(
19
+ load_in_8bit=True, # Использование 8-битного квантования
20
+ llm_int8_threshold=6.0, # Порог активации для 8-битных весов
21
+ )
22
+
23
  model = AutoModelForCausalLM.from_pretrained(
24
  "IlyaGusev/saiga_gemma2_10b",
25
  token=hf_token,
26
  torch_dtype=torch.float16, # Использование float16 для уменьшения потребления памяти
27
+ device_map="auto", # Автоматическое распределение модели на GPU
28
+ quantization_config=quant_config, # Применение конфигурации квантования
29
  )
30
 
 
31
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
32
  model.to(device)
33
 
 
52
  except Exception as e:
53
  return f"Извините, произошла ошибка при генерации ответа: {str(e)}"
54
 
 
55
  init_model()
56
 
 
57
  iface = gr.Interface(
58
  fn=generate_response,
59
  inputs=[
60
+ gr.Textbox(lines=2, placeholder="Введите ваш текст здесь..."),
61
+ gr.Slider(20, 200, step=1, default=100, label="Максимальная длина"),
62
+ gr.Slider(0.1, 1.0, step=0.1, default=0.7, label="Температура"),
63
+ gr.Slider(0.1, 1.0, step=0.05, default=0.85, label="Top-p"),
64
+ gr.Slider(1.0, 2.0, step=0.1, default=1.1, label="Штраф за повторение")
65
  ],
66
  outputs="text",
67
  title="LLM Model Demo",