import gradio as gr
import json
import numpy as np
import faiss
from transformers import LlamaTokenizer, LlamaModel
import torch

# Load model and tokenizer
model_name = "openlm-research/open_llama_3b_v2"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = LlamaModel.from_pretrained(model_name)
model.eval()

# Generate embeddings function
def generate_embeddings():
    with open("product.json", "r") as f:
        products = json.load(f)

    embeddings = []
    for product in products:
        inputs = tokenizer(product["description"], return_tensors="pt", truncation=True, max_length=512)
        with torch.no_grad():
            outputs = model(**inputs, output_hidden_states=True)
        embedding = outputs.hidden_states[-1].mean(dim=1).squeeze().numpy()
        embeddings.append({"id": product["id"], "embedding": embedding})

    dimension = embeddings[0]["embedding"].shape[0]
    index = faiss.IndexFlatL2(dimension)
    embedding_matrix = np.array([e["embedding"] for e in embeddings])
    index.add(embedding_matrix)
    faiss.write_index(index, "product_index.faiss")
    np.save("product_ids.npy", np.array([e["id"] for e in embeddings]))
    return "Embeddings generated successfully!"

# Find similar products function
def find_similar(product_id):
    index = faiss.read_index("product_index.faiss")
    product_ids = np.load("product_ids.npy")
    with open("product.json", "r") as f:
        products = {p["id"]: p for p in json.load(f)}

    product_index = np.where(product_ids == int(product_id))[0]
    if len(product_index) == 0:
        return "Product ID not found"
    
    query_embedding = index.reconstruct(product_index[0])
    distances, indices = index.search(np.expand_dims(query_embedding, axis=0), k=5)
    similar_ids = product_ids[indices[0]].tolist()
    similar_products = [products[pid]["description"] for pid in similar_ids if pid != int(product_id)]
    return similar_products

# Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("### Product Embedding and Similarity Search")
    with gr.Row():
        generate_button = gr.Button("Generate Embeddings")
        generate_output = gr.Textbox(label="Status")
        generate_button.click(generate_embeddings, outputs=generate_output)

    with gr.Row():
        product_id_input = gr.Textbox(label="Enter Product ID")
        similar_products_output = gr.Textbox(label="Similar Products")
        find_button = gr.Button("Find Similar Products")
        find_button.click(find_similar, inputs=product_id_input, outputs=similar_products_output)

# Enable public access
demo.launch(server_name="0.0.0.0", server_port=7860, share=True)