Update app.py

app.py

@@ -4,28 +4,23 @@ import pandas as pd
 import numpy as np
 from ast import literal_eval
 
-# Dropdown to select the model
+
+st.title("Mining Potential Legitimate Domains from a Typosquatted Domain")
 model_choice = st.selectbox("Select the embedding model:", ["", "Embedder-typosquat-detect-Canine", "Embedder-typosquat-detect"], index=0)
 
-# Load the model only if a model is selected
 if model_choice:
     model = SentenceTransformer(f"./{model_choice}")
 
-    # Load the domains and embeddings
     domains_df = pd.read_csv(f'./{model_choice}/domains_embs.csv')
     domains_df.embedding = domains_df.embedding.apply(literal_eval)
     corpus_domains = domains_df.domain.to_list()
     corpus_embeddings = np.stack(domains_df.embedding.values).astype(np.float32)  # Ensure embeddings are float32
 
-    # Streamlit App
-    st.title("Mining Potential Legitimate Domains from a Typosquatted Domain")
     st.write("Enter a potential typosquatted domain and select the number of top results to retrieve.")
 
-    # User Inputs
     domain = st.text_input("Potential Typosquatted Domain")
     top_k = st.number_input("Top K Results", min_value=1, max_value=50, value=5, step=1)
 
-    # Button to trigger search
     if st.button("Search for Legitimate Domains"):
         if domain:
             # Perform Semantic Search
@@ -33,12 +28,10 @@ if model_choice:
             semantic_res = util.semantic_search(query_emb, corpus_embeddings, top_k=top_k)[0]
             ids = [r['corpus_id'] for r in semantic_res]
             scores = [r['score'] for r in semantic_res]
-            
-            # Create a DataFrame for the results
+
             res_df = domains_df.loc[ids, ['domain']].copy()
             res_df['score'] = scores
             
-            # Display the result DataFrame
             st.write("Mined Domains:")
             st.dataframe(res_df)
         else: