heimoshuiyu/local-embedding-api
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: cpu模式不支持 float16

Browse Source
This commit is contained in:
黑墨水鱼
2024-06-04 17:57:15 +08:00
parent 947cc86ad8
commit 41c1205dfb
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
acge_embedding.py
View File

@@ -6,10 +6,10 @@ import torch
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device) print("Using device:", device)
model_name = "models/aspire--acge-large-text" model_name = "models/aspire--acge-large-zh"
print("Loading model", model_name) print("Loading model", model_name)
model = ( model = (
AutoModel.from_pretrained(model_name, local_files_only=True, torch_dtype=torch.float16).eval().to(device) AutoModel.from_pretrained(model_name, local_files_only=True).eval().to(device)
) )
tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True) tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
print("Model", model_name, "loaded!") print("Model", model_name, "loaded!")