chen666-666 commited on
Commit
0378c00
·
1 Parent(s): 4fab3f7

add app.py and requirements.txt

Browse files
Files changed (1) hide show
  1. app.py +9 -8
app.py CHANGED
@@ -17,20 +17,21 @@ bert_ner_pipeline = pipeline("ner", model=bert_ner_model, tokenizer=bert_tokeniz
17
  chatglm_model, chatglm_tokenizer = None, None
18
  use_chatglm = False
19
  try:
20
- # 强制使用CPU加载ChatGLM
21
- chatglm_model_name = "THUDM/chatglm3-6b"
22
- chatglm_tokenizer = AutoTokenizer.from_pretrained(chatglm_model_name, trust_remote_code=True)
 
 
23
  chatglm_model = AutoModel.from_pretrained(
24
  chatglm_model_name,
25
  trust_remote_code=True,
26
- device_map="cpu", # 强制使用CPU
27
- torch_dtype=torch.float32, # 使用FP32精度
28
- low_cpu_mem_usage=True # 优化内存使用
29
  ).eval()
30
  use_chatglm = True
31
- print("✅ ChatGLM 已加载到 CPU")
32
  except Exception as e:
33
- print(f"❌ ChatGLM 加载失败: {e}")
34
 
35
  # ======================== 知识图谱结构 ========================
36
  knowledge_graph = {"entities": set(), "relations": set()}
 
17
  chatglm_model, chatglm_tokenizer = None, None
18
  use_chatglm = False
19
  try:
20
+ chatglm_model_name = "THUDM/chatglm-6b-int4" # 4-bit量化版本
21
+ chatglm_tokenizer = AutoTokenizer.from_pretrained(
22
+ chatglm_model_name,
23
+ trust_remote_code=True
24
+ )
25
  chatglm_model = AutoModel.from_pretrained(
26
  chatglm_model_name,
27
  trust_remote_code=True,
28
+ device_map="cpu",
29
+ torch_dtype=torch.float32 # 必须使用float32
 
30
  ).eval()
31
  use_chatglm = True
32
+ print("✅ 4-bit量化版ChatGLM加载成功(需6GB内存)")
33
  except Exception as e:
34
+ print(f"❌ 量化模型加载失败: {e}")
35
 
36
  # ======================== 知识图谱结构 ========================
37
  knowledge_graph = {"entities": set(), "relations": set()}