Update app.py

app.py

@@ -1,43 +1,56 @@
 import gradio as gr
 from sentence_transformers import SentenceTransformer, util
-import ast
+import re
 
 # Load the SentenceTransformer model
 model = SentenceTransformer('all-MiniLM-L6-v2')
 
-def compare_embeddings(query, text_lists):
-    # Convert string representation of lists to actual lists
-    try:
-        lists = ast.literal_eval(text_lists)
-    except:
-        return "Error: Invalid input format for text lists."
+def extract_lists(text):
+    # Use regex to find all lists in the format ['item1', 'item2', ...]
+    pattern = r'\[([^\]]+)\]'
+    matches = re.findall(pattern, text)
+    
+    # Process each match into a list of keywords
+    lists = []
+    for match in matches:
+        # Split by comma, strip whitespace and quotes, and filter out empty strings
+        keywords = [item.strip().strip("'\"") for item in match.split(',') if item.strip()]
+        if keywords:  # Only add non-empty lists
+            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, lst in enumerate(lists):
-        # Encode the list items
-        list_embeddings = model.encode(lst, convert_to_tensor=True)
+    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, list_embeddings)[0]
+        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))
+        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 in results:
+    for i, score, keywords in results:
         output += f"List {i}: Similarity score {score:.4f}\n"
-        for j, item in enumerate(lists[i]):
-            output += f"  Element {j}: {item}\n"
-        output += "\n"
+        output += f"  Keywords: {', '.join(keywords)}\n\n"
 
     return output
 
@@ -46,11 +59,15 @@ iface = gr.Interface(
     fn=compare_embeddings,
     inputs=[
         gr.Textbox(label="Query"),
-        gr.Textbox(label="Text Lists (in format [['', ''], ['', ''], ['', '']])"),
+        gr.Textbox(
+            label="Lists of keywords",
+            placeholder="Enter lists in the format: ['keyword1', 'keyword2', ...] ['keyword3', 'keyword4', ...]",
+            lines=5
+        ),
     ],
     outputs=gr.Textbox(label="Results"),
-    title="Embedding Comparison App",
-    description="Compare a query with multiple lists of text and find the most relevant list."
+    title="Keyword Lists Comparison App",
+    description="Compare a query with multiple lists of keywords and find the most relevant lists. Enter each list in square brackets, separated by commas."
 )
 
 # Launch the app