import gradio as gr
from sentence_transformers import SentenceTransformer, util
import re

# Load the SentenceTransformer model
model = SentenceTransformer('all-MiniLM-L6-v2')

def extract_lists(text):
    # Split the input text by newlines and process each line
    lines = text.split('\n')
    lists = []
    for line in lines:
        # Strip whitespace and check if the line is not empty
        line = line.strip()
        if line:
            # Split the line by spaces and strip any remaining whitespace
            keywords = [item.strip() for item in line.split() if item.strip()]
            if keywords:
                lists.append(keywords)
    return lists

def compare_embeddings(query, lists_text):
    # Extract lists from the input text
    keyword_lists = extract_lists(lists_text)
    if not keyword_lists:
        return "No valid lists found in the input. Please check the format."

    # Encode the query
    query_embedding = model.encode(query, convert_to_tensor=True)

    results = []
    for i, keywords in enumerate(keyword_lists):
        # Encode the keywords
        keyword_embeddings = model.encode(keywords, convert_to_tensor=True)

        # Calculate cosine similarity
        similarities = util.pytorch_cos_sim(query_embedding, keyword_embeddings)[0]

        # Calculate average similarity for the list
        avg_similarity = similarities.mean().item()
        results.append((i, avg_similarity, keywords))

    # Sort results by similarity score (descending)
    results.sort(key=lambda x: x[1], reverse=True)

    # Format the output
    output = ""
    for i, score, keywords in results:
        output += f"List {i}: Similarity score {score:.4f}\n"
        output += f" Keywords: {' '.join(keywords)}\n\n"

    return output

# Create the Gradio interface
iface = gr.Interface(
    fn=compare_embeddings,
    inputs=[
        gr.Textbox(label="Query"),
        gr.Textbox(
            label="Lists of keywords",
            placeholder="Enter each list of keywords on a new line, with keywords separated by spaces.",
            lines=5
        ),
    ],
    outputs=gr.Textbox(label="Results"),
    title="Keyword Lists Comparison App",
    description="Compare a query with multiple lists of keywords and find the most relevant lists. Enter each list on a new line, with keywords separated by spaces."
)

# Launch the app
iface.launch()