Spaces:
Paused
Paused
File size: 1,426 Bytes
2400080 | 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 42 | import json
import numpy as np
import faiss
from transformers import LlamaTokenizer, LlamaModel
import torch
# Load the model and tokenizer
model_name = "openlm-research/open_llama_3b"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = LlamaModel.from_pretrained(model_name)
model.eval()
# Generate embeddings function
def get_embedding(text, model, tokenizer):
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
outputs = model(**inputs, output_hidden_states=True)
# Use the mean of the last hidden state as the embedding
return outputs.hidden_states[-1].mean(dim=1).squeeze().numpy()
# Load products from JSON
with open("products.json", "r") as f:
products = json.load(f)
# Generate embeddings
product_embeddings = []
for product in products:
embedding = get_embedding(product["description"], model, tokenizer)
product_embeddings.append({"id": product["id"], "embedding": embedding})
# Save embeddings to FAISS
dimension = product_embeddings[0]["embedding"].shape[0]
index = faiss.IndexFlatL2(dimension)
embedding_matrix = np.array([p["embedding"] for p in product_embeddings])
index.add(embedding_matrix)
# Save FAISS index and product IDs
faiss.write_index(index, "product_index.faiss")
np.save("product_ids.npy", np.array([p["id"] for p in product_embeddings]))
print("Embeddings generated and saved successfully!")
|