import gradio as gr
import fasttext
from huggingface_hub import hf_hub_download
from huggingface_hub import login
import os

login(token=os.environ["HF_TOKEN"])


model_path = hf_hub_download(repo_id="LuminAISecurity/fasttext-english-nearest-neighbors", filename="model.bin")
model = fasttext.load_model(model_path)


def get_nearest_neighbors(word, model, k=5):
    neighbors = model.get_nearest_neighbors(word, k=k)
    return neighbors


# Define the K-nearest neighbors function
def find_nearest_neighbors(text, k):
    neighbors = model.get_nearest_neighbors(text, k=k+3)
    return neighbors[2:]


# Create the Gradio interface
iface = gr.Interface(
    fn=find_nearest_neighbors,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter text here...", label="Input Text"),
        gr.Slider(1, 10, value=5, label="Number of Neighbors")
    ],
    outputs=gr.JSON(label="Nearest Neighbors"),
    title="FastText K-Nearest Neighbors Finder",
    description="Enter a text and the number of neighbors to find the closest matches using the FastText model.",
    examples=[["king", 5], ["queen", 3]],
    theme="default",
    css="style.css"
)

# Launch the app
if __name__ == "__main__":
    iface.launch(share=True)