import gradio as gr
from wordllama import WordLlama

# Load the default WordLlama model
wl = WordLlama.load()

# Define the function that ranks candidate documents based on their similarity to the query
def rank_documents(query, candidates):
    candidates_list = candidates.split("\n")  # Split candidates by new lines
    ranked_docs = wl.rank(query, candidates_list)
    return ranked_docs

# Define the Gradio interface with examples included
iface = gr.Interface(
    fn=rank_documents,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your query sentence..."),
        gr.Textbox(lines=5, placeholder="Enter candidate sentences (one per line)...")
    ],
    outputs=gr.Textbox(label="Ranked Documents"),
    title="Document Ranking with WordLlama",
    description="Rank candidate sentences based on their similarity to the query using WordLlama.",
    examples=[
        ["i went to the car", "i went to the park\ni went to the shop\ni went to the truck\ni went to the vehicle"],
        ["I love pizza", "I enjoy eating pasta\nPizza is my favorite food\nI love burgers\nPasta is tasty too"],
        ["The sky is blue", "The sky is clear\nThe ocean is deep\nThe sky is cloudy\nThe sky is bright"],
        ["She is reading a book", "She loves reading\nShe is studying\nHe is playing\nShe enjoys books"],
        ["He went to the gym", "She is going to the park\nHe went to the store\nHe is working out\nHe went home"]
    ]
)

iface.launch()