chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,41 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagPseudoMoEModel
|
||||
|
||||
|
||||
def test_pseudo_moe_single_device():
|
||||
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="coding",
|
||||
use_fp16=False,
|
||||
use_bf16=True,
|
||||
trust_remote_code=True,
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv("HF_HUB_CACHE", None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define",
|
||||
] * 10
|
||||
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.",
|
||||
] * 10
|
||||
|
||||
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_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.700 0.246]\n [0.158 0.654]]")
|
||||
Reference in New Issue
Block a user