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

lfs add acge-large-text

Browse Source
This commit is contained in:
黑墨水鱼
2024-06-04 12:29:22 +08:00
parent dea0ba42ff
commit 4201ad34ff
6 changed files with 18 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
acge_embedding.py
View File

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