several imporvements

app.py

@@ -10,9 +10,16 @@ required_columns = [
     'Avg forward segment size', 'Avg backward segment size'
 ]
 
+# Hugging Face API Settings
+HUGGING_FACE_API_URL = "https://api-inference.huggingface.co/models/sentence-transformers/all-distilroberta-v1"
+
 # Streamlit UI
 st.title("NetFlow Log Comparison Tool")
-st.write("Compare your NetFlow logs against Sigma rules or MITRE ATT&CK patterns using RAG.")
+st.write("Compare your NetFlow logs against Sigma rules or MITRE ATT&CK patterns using Retrieval-Augmented Generation (RAG).")
+
+# Display the embedding model being used
+st.write("### Embedding Model in Use")
+st.write("The model used for embedding is: **All-DistilRoBERTa-V1**")
 
 # Instructions for data upload
 st.markdown("""
@@ -22,9 +29,17 @@ st.markdown("""
 - You can upload **up to 5 rows** for analysis.
 """)
 
-# Display required schema for users
+# Display required schema for users with bullet points
 st.write("### Required NetFlow Schema:")
-st.write(", ".join(required_columns))
+st.markdown("""
+- **Flow duration**
+- **Source port**
+- **Destination port**
+- **Total forward packets**
+- **Total backward packets**
+- **Avg forward segment size**
+- **Avg backward segment size**
+""")
 
 # Step 1: File Upload
 uploaded_file = st.file_uploader("Upload your NetFlow log sequence CSV file", type="csv")
@@ -34,7 +49,12 @@ hugging_face_api_token = st.text_input("Enter your Hugging Face API Token", type
 if not hugging_face_api_token:
     st.warning("Please provide a Hugging Face API Token to proceed.")
 
-# Step 3: Run Comparison if File Uploaded and Token Provided
+# Step 3: Model and Comparison Options
+st.write("### Model and Comparison Options")
+llm_choice = st.selectbox("Select LLM", ["All-DistilRoBERTa-V1"])  # Add other models as necessary
+comparison_choice = st.selectbox("Select Comparison Type", ["Mitre", "Sigma"])
+
+# Step 4: Run Comparison if File Uploaded and Token Provided
 if uploaded_file and hugging_face_api_token:
     # Read and display the file using CSV module
     csv_file = StringIO(uploaded_file.getvalue().decode("utf-8"))
@@ -54,8 +74,7 @@ if uploaded_file and hugging_face_api_token:
             # Prepare data for Hugging Face API call
             input_texts = [f"{row}" for row in csv_data]  # Convert each row to a string for comparison
             
-            # Step 4: Call Hugging Face API
-            HUGGING_FACE_API_URL = "https://api-inference.huggingface.co/models/sentence-transformers/all-distilroberta-v1"
+            # Call Hugging Face API
             headers = {"Authorization": f"Bearer {hugging_face_api_token}"}
 
             try:
@@ -66,7 +85,14 @@ if uploaded_file and hugging_face_api_token:
                 # Display the results
                 st.write("### Comparison Results")
                 comparison_results = response.json()
-                st.write(comparison_results)
+
+                # Sort and extract top 3 results for display
+                top_results = sorted(comparison_results, key=lambda x: x['score'], reverse=True)[:3]
+                
+                # Display the top 3 results
+                for idx, result in enumerate(top_results):
+                    st.write(f"**{idx + 1}.** Matched Sequence: `{result['sequence']}`")
+                    st.write(f"   - **Cosine Similarity Score**: {result['score']:.4f}")
 
             except requests.exceptions.RequestException as e:
                 st.error(f"Error calling Hugging Face API: {str(e)}")
@@ -84,7 +110,7 @@ st.write("We value your feedback. [Fill out our survey](https://docs.google.com/
 # Footer
 st.markdown("---")
 st.write("This free site is maintained by DeepTempo.")
-# st.image("Final_DeepTempo_logo.png", width=300) # Adjust the path and width as needed 'Final DeepTempo logo.png
+# st.image("Final_DeepTempo_logo.png", width=300) # Adjust the path and width as needed 'Final DeepTempo logo.png'
 st.write("[Visit DeepTempo.ai](https://deeptempo.ai)")
 st.write("[Check out the underlying code on GitHub](https://github.com/deepsecoss)")