簡嘉琳 commited on
Commit
705c0ca
·
1 Parent(s): 5f37119

Fix trust_remote_code issue

Browse files
Files changed (1) hide show
  1. app.py +16 -14
app.py CHANGED
@@ -5,43 +5,45 @@ import torch
5
  # 預先定義 Hugging Face 模型
6
  MODEL_NAMES = {
7
  "DeepSeek-V3": "deepseek-ai/DeepSeek-V3",
8
- "DeepSeek-R1": "deepseek-ai/DeepSeek-R1"
9
  }
10
 
11
- def load_model(model_name):
12
- """載入 Hugging Face 模型與 tokenizer"""
13
- model_path = MODEL_NAMES.get(model_name, "deepseek-ai/DeepSeek-V3")
14
- tokenizer = AutoTokenizer.from_pretrained(model_path)
15
- model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16).cuda()
 
16
  return model, tokenizer
17
 
 
18
  # 預設載入 DeepSeek-V3
19
  current_model, current_tokenizer = load_model("DeepSeek-V3")
20
 
 
21
  def chat(message, history, model_name):
22
  """處理聊天訊息"""
23
  global current_model, current_tokenizer
24
-
25
  # 若模型不同則切換
26
  if model_name != current_model:
27
  current_model, current_tokenizer = load_model(model_name)
28
-
29
  inputs = current_tokenizer(message, return_tensors="pt").to("cuda")
30
  outputs = current_model.generate(**inputs, max_length=1024)
31
  response = current_tokenizer.decode(outputs[0], skip_special_tokens=True)
32
-
33
  return response
34
 
 
35
  with gr.Blocks() as app:
36
  gr.Markdown("## Chatbot with DeepSeek Models")
37
-
38
  with gr.Row():
39
  chat_interface = gr.ChatInterface(chat, streaming=True, save_history=True)
40
  model_selector = gr.Dropdown(
41
- choices=list(MODEL_NAMES.keys()),
42
- value="DeepSeek-V3",
43
- label="Select Model"
44
  )
45
-
46
  chat_interface.append(model_selector)
47
  app.launch()
 
5
  # 預先定義 Hugging Face 模型
6
  MODEL_NAMES = {
7
  "DeepSeek-V3": "deepseek-ai/DeepSeek-V3",
8
+ "DeepSeek-R1": "deepseek-ai/DeepSeek-R1",
9
  }
10
 
11
+
12
+ def load_model(model_path):
13
+ tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
14
+ model = AutoModelForCausalLM.from_pretrained(
15
+ model_path, torch_dtype=torch.float16, trust_remote_code=True
16
+ ).cuda()
17
  return model, tokenizer
18
 
19
+
20
  # 預設載入 DeepSeek-V3
21
  current_model, current_tokenizer = load_model("DeepSeek-V3")
22
 
23
+
24
  def chat(message, history, model_name):
25
  """處理聊天訊息"""
26
  global current_model, current_tokenizer
27
+
28
  # 若模型不同則切換
29
  if model_name != current_model:
30
  current_model, current_tokenizer = load_model(model_name)
31
+
32
  inputs = current_tokenizer(message, return_tensors="pt").to("cuda")
33
  outputs = current_model.generate(**inputs, max_length=1024)
34
  response = current_tokenizer.decode(outputs[0], skip_special_tokens=True)
35
+
36
  return response
37
 
38
+
39
  with gr.Blocks() as app:
40
  gr.Markdown("## Chatbot with DeepSeek Models")
41
+
42
  with gr.Row():
43
  chat_interface = gr.ChatInterface(chat, streaming=True, save_history=True)
44
  model_selector = gr.Dropdown(
45
+ choices=list(MODEL_NAMES.keys()), value="DeepSeek-V3", label="Select Model"
 
 
46
  )
47
+
48
  chat_interface.append(model_selector)
49
  app.launch()