Views
No views yet
txtai using the following code.1import txtai
2
3embeddings = txtai.Embeddings(
4 path="neuml/biomedbert-hash-nano-embeddings/biomedbert-hash-nano-embeddings-fp16.tflite"
5 content=True,
6)
7embeddings.index(documents())
8
9# Run a query
10embeddings.search("query to run")1#
2# pip install litert-torch txtai
3#
4# See https://github.com/google-ai-edge/litert-torch
5#
6
7import argparse
8import json
9import os
10
11import litert_torch
12import torch
13
14from ai_edge_quantizer import quantizer, recipe
15
16from litert_torch.generative.quantize import quant_recipes
17from torch import nn
18from transformers import AutoTokenizer
19from txtai.models import PoolingFactory
20from txtai.util import Download
21
22
23class Pooling(nn.Module):
24 def __init__(self, path, device, **kwargs):
25 super().__init__()
26 self.model = PoolingFactory.create({"path": path, "device": device, "modelargs": kwargs})
27
28 # Read max length parameter. Don't use tokenizer max length since model is exported as static shape
29 config = f"{path}/sentence_bert_config.json"
30 config = config if os.path.exists(config) else Download()(config)
31
32 with open(config, encoding="utf-8") as f:
33 data = json.load(f)
34 self.maxlength = data["max_seq_length"]
35
36 # pylint: disable=W0221
37 def forward(self, input_ids=None, attention_mask=None, token_type_ids=None):
38 inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
39 if token_type_ids is not None:
40 inputs["token_type_ids"] = token_type_ids
41
42 return self.model.forward(**inputs)
43
44
45def export(args):
46 model = Pooling(args.input, -1).float().eval()
47
48 batch, maxlength = 4, model.maxlength
49 inputs = (
50 torch.ones(batch, maxlength, dtype=torch.int32),
51 torch.ones(batch, maxlength, dtype=torch.int32),
52 torch.ones(batch, maxlength, dtype=torch.int32),
53 )
54
55 base = os.path.basename(args.input)
56
57 if args.quant == "int8":
58 config, path = quant_recipes.full_dynamic_recipe(), f"{base}-int8.tflite"
59 elif args.quant == "fp16":
60 config, path = quant_recipes.full_fp16_recipe(), f"{base}-fp16.tflite"
61 else:
62 config, path = None, f"{base}-fp32.tflite"
63
64 # Create output directory
65 os.makedirs(args.output, exist_ok=True)
66 path = os.path.join(args.output, path)
67
68 # Save tokenizer
69 tokenizer = AutoTokenizer.from_pretrained(args.input)
70 tokenizer.save_pretrained(args.output)
71
72 # Save model
73 model = litert_torch.convert(model, inputs, quant_config=config)
74 model.export(path)
75
76 # Quantize to int4, if necessary
77 if args.quant == "int4":
78 qt = quantizer.Quantizer(path, recipe.dynamic_wi4_afp32())
79
80 path = os.path.join(args.output, f"{base}-int4.tflite")
81 qt.quantize().export_model(path, overwrite=True)
82
83 return path
84
85
86if __name__ == "__main__":
87 parser = argparse.ArgumentParser()
88 parser.add_argument("--input", help="model path", required=True)
89 parser.add_argument("--output", help="model output directory", required=True)
90 parser.add_argument("--quant", help="model quantization", choices=["int4", "int8", "fp16"])
91 args = parser.parse_args()
92
93 export(args)