Update app.py

app.py

@@ -2,7 +2,7 @@ from sklearn.preprocessing import OrdinalEncoder
 import xgboost as xgb
 import numpy as np
 import matplotlib.pyplot as plt
-import argparse, h5py, os, re
+import argparse, h5py, os, re, zipfile
 import streamlit as st
 
 def get_POMFinder():
@@ -96,6 +96,36 @@ def download_button(file_name, button_text):
         file_name=file_name,
         mime="text/xyz",)
 
+def create_zip_with_txt(y_onehotenc_cat, res, y_pred_proba):
+    """
+    Create a zip file containing specific .xyz files and a txt file with model guesses.
+    
+    Parameters:
+        y_onehotenc_cat: The one-hot encoder categories object.
+        res: The result index for selecting the appropriate file.
+        y_pred_proba: The probability associated with each prediction.
+        
+    Returns:
+        None. A zip file is created in the current directory.
+    """
+    
+    zip_file_name = "POMFinder_results.zip"
+    
+    # Create a zip file
+    with zipfile.ZipFile(zip_file_name, 'w') as zipf:
+        # Add the .xyz files
+        for i in range(5):
+            file_path = "Backend/COD_ICSD_XYZs_POMs_unique99/" + str(y_onehotenc_cat.categories_[i][res[0]])[2:-2] + "cale.xyz"
+            zipf.write(file_path, arcname=str(y_onehotenc_cat.categories_[i][res[0]])[2:-2] + "cale.xyz")
+        
+        # Create and add the txt file with model guesses
+        with zipf.open("model_guesses.txt", "w") as f:
+            for i in range(5):
+                guess_text = f'The {i+1}st guess from the model is: {str(y_onehotenc_cat.categories_[i][res[0]])[2:-2]+"cale.xyz"} with a probability of {y_pred_proba[res[0]]:.2f} %\n'
+                f.write(guess_text.encode('utf-8'))
+
+
+
 
 st.title('POMFinder')
 st.write('Welcome to POMFinder which is a tree-based supervised learning algorithm that can predict the polyoxometalate cluster from a Pair Distribution Function.')
@@ -136,7 +166,8 @@ else:
 
     # Download the structural database
     #download_button("COD_ICSD_XYZs_POMs_unique99.zip", "Download structural database")
-    download_button("Backend/COD_ICSD_XYZs_POMs_unique99/"+str(y_onehotenc_cat.categories_[0][res[0]])[2:-2]+"cale.xyz", "Download top-5 predictions")
+    create_zip_with_txt(y_onehotenc_cat, res, y_pred_proba)
+    download_button("POMFinder_results.zip", "Download top-5 predictions")
 
 st.subheader('Cite')