File size: 1,383 Bytes
0908771 1e25956 0908771 1e25956 0908771 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """Run from the downloaded release: python example.py [--device cuda:0]."""
import argparse
import sys
from pathlib import Path
import numpy as np
# HF snapshots use symlinks; retain the snapshot directory for sibling imports.
sys.path.insert(0, str(Path(__file__).absolute().parent))
from qime import QIMEModel
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--device", default=None)
parser.add_argument("--local-files-only", action="store_true")
args = parser.parse_args()
model = QIMEModel.from_pretrained(
Path(__file__).parent, device=args.device, local_files_only=args.local_files_only,
)
texts = [
"Patient presents with severe chest pain.",
"Treatment involves daily insulin injections.",
]
embeddings = model.encode(texts)
assert embeddings.shape == (2, 8855)
assert np.all(embeddings.sum(axis=1) == 256)
assert np.isin(embeddings, [0, 1]).all()
print("Embedding shape:", embeddings.shape)
print("Active dimensions per text:", embeddings.sum(axis=1))
for text, embedding in zip(texts, embeddings):
print("\nText:", text)
print("First 10 active coordinates (coordinate order, not relevance order):")
for index in np.flatnonzero(embedding)[:10]:
print(f"[{index}] {model.questions[index]}")
if __name__ == "__main__":
main()
|