import json
import numpy as np
import faiss
import gradio as gr
from transformers import AutoTokenizer, TFAutoModel

# Load configuration
with open("config.json", "r") as f:
    config = json.load(f)

# Load model and tokenizer
model_ckpt = config["model_checkpoint"]
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = TFAutoModel.from_pretrained(model_ckpt, from_pt=True)

# Load embeddings and URLs
embeddings = np.load(config["embedding_path"])
urls = np.load(config["urls_path"])
texts = np.load(config["texts_path"])

# Load FAISS index
index = faiss.read_index(config["faiss_index_path"])

# Query embedding function for search
def get_embeddings_batch(texts, batch_size=config["batch_size"]):
    embeddings = []
    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i:i + batch_size]
        inputs = tokenizer(batch_texts, padding=True, truncation=True, return_tensors='tf')
        outputs = model(**inputs)
        batch_embeddings = outputs.last_hidden_state[:, 0, :].numpy()  # Use CLS token
        embeddings.append(batch_embeddings)
    return np.vstack(embeddings)

def semantic_search(query):
    query_embedding = get_embeddings_batch([query])
    D, I = index.search(query_embedding, config["num_results"])
    similar_texts = [texts[i] for i in I[0]]
    similar_urls = [urls[i] for i in I[0]]
    return {"texts": similar_texts, "urls": similar_urls, "distances": D.tolist()}

def query_interface(query):
    result = semantic_search(query)
    return result

iface = gr.Interface(
    fn=query_interface,
    inputs=gr.components.Textbox(lines=2, placeholder="Enter your query here..."),
    outputs=gr.components.JSON(),
    title="<b>Windsurfing Technique Search</b>",
    description="Discover core information on windsurfing techniques. "
)

if __name__ == "__main__":
    iface.launch()