import gradio as gr
import time
from sentence_transformers import CrossEncoder

# Load model
model2 = CrossEncoder('enochlev/coherence-all-mpnet-base-v2')

# Predefined examples
examples = [
    ["What is your favorite color?", "Blue!"],
    ["Do you like playing outside?", "I like ice cream."],
    ["What is your favorite animal?", "I like dogs!"],
    ["Do you want to go to the park?", "Yes, I want to go on the swings!"],
    ["What is your favorite food?", "I like playing with blocks."],
    ["Do you have a pet?", "Yes, I have a cat named Whiskers."],
    ["What is your favorite thing to do on a sunny day?", "I like playing soccer with my friends."]
]

# Global index for cycling through examples
current_index = 0

def check_coherence(sentence1: str, sentence2: str) -> float:
    """
    Predicts the coherence score for a pair of sentences.
    """
    score = model2.predict([[sentence1, sentence2]])[0]
    return score

def get_next_example():
    """
    Returns the next example pair and updates the current index.
    """
    global current_index
    pair = examples[current_index]
    current_index = (current_index + 1) % len(examples)
    return pair[0], pair[1]

with gr.Blocks() as demo:
    gr.Markdown("## Coherence Checker Demo")
    
    with gr.Row():
        with gr.Column():
            inp1 = gr.Textbox(label="Sentence 1", placeholder="Enter first sentence...")
            inp2 = gr.Textbox(label="Sentence 2", placeholder="Enter second sentence...")
            check_button = gr.Button("Check Coherence")
            next_example_button = gr.Button("Next Example")
        with gr.Column():
            output_score = gr.Textbox(label="Coherence Score", interactive=False)
    
    check_button.click(fn=check_coherence, inputs=[inp1, inp2], outputs=output_score)
    next_example_button.click(fn=get_next_example, inputs=[], outputs=[inp1, inp2])
    
demo.launch()