| """CPU-only BGE subprocess. No spaces import or GPU model initialization.""" |
| import argparse |
| import json |
| import sys |
| import torch |
| import torch.nn.functional as F |
| from transformers import AutoTokenizer,AutoModel |
|
|
|
|
| def main(): |
| parser=argparse.ArgumentParser() |
| parser.add_argument('--source',required=True) |
| args=parser.parse_args() |
| torch.set_num_threads(min(4,torch.get_num_threads())) |
| tokenizer=AutoTokenizer.from_pretrained(args.source,local_files_only=True) |
| model=AutoModel.from_pretrained(args.source,local_files_only=True,attn_implementation='eager').eval().to('cpu') |
| for line in sys.stdin: |
| try: |
| texts=json.loads(line)['texts'] |
| if len(texts)>128 or sum(map(len,texts))>160000: |
| raise ValueError('Embedding request exceeds bounds') |
| vectors=[] |
| with torch.inference_mode(): |
| for start in range(0,len(texts),16): |
| tokens=tokenizer(texts[start:start+16],padding=True,truncation=True,max_length=512,return_tensors='pt') |
| values=model(**tokens).last_hidden_state[:,0] |
| vectors.extend(F.normalize(values.float(),dim=-1).tolist()) |
| print(json.dumps({'vectors':vectors}),flush=True) |
| except Exception as exc: |
| print(json.dumps({'error':type(exc).__name__}),flush=True) |
|
|
|
|
| if __name__=='__main__': |
| try: |
| main() |
| except Exception as exc: |
| print(json.dumps({'error':'embedding_startup_'+type(exc).__name__}),flush=True) |
|
|
