wehub-resource-sync
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import os
|
|
from FlagEmbedding import FlagPseudoMoEModel
|
|
|
|
|
|
def test_pseudo_moe_multi_devices():
|
|
model_name_or_path = "geevec-ai/geevec-embeddings-1.0-lite"
|
|
|
|
model = FlagPseudoMoEModel(
|
|
model_name_or_path,
|
|
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
|
query_instruction_format="Instruct: {}\nQuery: {}",
|
|
domain_for_pseudo_moe="reasoning",
|
|
use_fp16=False,
|
|
use_bf16=True,
|
|
trust_remote_code=True,
|
|
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
|
cache_dir=os.getenv("HF_HUB_CACHE", None),
|
|
)
|
|
|
|
queries = [
|
|
"how much protein should a female eat",
|
|
"summit define",
|
|
] * 100
|
|
passages = [
|
|
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day.",
|
|
"Definition of summit for English Language Learners: the highest point of a mountain; the highest level; a meeting between leaders.",
|
|
] * 100
|
|
|
|
queries_embeddings = model.encode_queries(queries)
|
|
passages_embeddings = model.encode_corpus(passages)
|
|
|
|
cos_scores = queries_embeddings @ passages_embeddings.T
|
|
print(cos_scores[:2, :2])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_pseudo_moe_multi_devices()
|
|
|
|
print("--------------------------------")
|
|
print("Expected Output:")
|
|
print("[[0.844 0.466 ]\n [0.395 0.684 ]]")
|