37 lines
1006 B
Python
37 lines
1006 B
Python
import sys
|
|
import json
|
|
from core import embed_room as ap
|
|
from config.embed import MODEL, N_THREADS, N_CTX, DEFAULT_PREFIX
|
|
from llama_cpp import Llama
|
|
|
|
def embed_run(args):
|
|
|
|
if not MODEL.is_file():
|
|
print(f"Model not found: {MODEL}. Download it first (see README).", file=sys.stderr)
|
|
return
|
|
|
|
prefix = args.prefix if args.prefix is not None else DEFAULT_PREFIX
|
|
|
|
print("Loading model...")
|
|
llm = Llama(
|
|
model_path = str(MODEL),
|
|
n_ctx = N_CTX,
|
|
n_threads = N_THREADS,
|
|
embedding = True,
|
|
verbose = False,
|
|
)
|
|
print("Model ready!")
|
|
|
|
inputs = [prefix + t for t in args.texts]
|
|
|
|
resp = llm.create_embedding(inputs)
|
|
vectors = [d["embedding"] for d in resp["data"]]
|
|
|
|
out = [
|
|
{"text": text, "dim": len(vec), "embedding": vec}
|
|
for text, vec in zip(args.texts, vectors)
|
|
]
|
|
print(json.dumps(out, ensure_ascii=False, indent=2))
|
|
|
|
if __name__ == "__main__":
|
|
embed_run( ap.parser.parse_args() ) |