Views
No views yet
| Attribute | Value |
|---|---|
| Task | Sentence Embeddings |
| Format | .tflite (Float32) |
| File Size | 86.3 MB |
| Input Length | 512 tokens |
| Output Dim | 384 |
1import numpy as np
2from ai_edge_litert.interpreter import Interpreter
3from transformers import AutoTokenizer
4
5# Load model
6interpreter = Interpreter(model_path="Snowflake_snowflake-arctic-embed-xs.tflite")
7interpreter.allocate_tensors()
8
9tokenizer = AutoTokenizer.from_pretrained("Snowflake/snowflake-arctic-embed-xs")
10
11def get_embedding(text):
12 inputs = tokenizer(text, max_length=512, padding="max_length", truncation=True, return_tensors="np")
13
14 input_details = interpreter.get_input_details()
15 interpreter.set_tensor(input_details[0]['index'], inputs['input_ids'].astype(np.int64))
16 interpreter.set_tensor(input_details[1]['index'], inputs['attention_mask'].astype(np.int64))
17
18 interpreter.invoke()
19
20 output_details = interpreter.get_output_details()
21 return interpreter.get_tensor(output_details[0]['index'])[0]
22
23emb = get_embedding("Hello world")
24print(f"Embedding shape: {emb.shape}")