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!")