Update app.py

app.py

@@ -25,8 +25,8 @@ class CFG():
     input_data = st.text_area(display_text)
     model_name_or_path = 'sagawa/ZINC-t5-productpredicition'
     model = 't5'
-    num_beams = 5
-    num_return_sequences = 5
+    num_beams = st.number_input(label='num beams', min_value=1, max_value=10, value=5, step=1)
+    num_return_sequences = st.number_input(label='num return sequences', min_value=1, max_value=num_beams, value=num_beams, step=1)
     seed = 42
 
 
@@ -79,7 +79,8 @@ if CFG.uploaded_file is not None:
     def convert_df(df):
         # IMPORTANT: Cache the conversion to prevent computation on every rerun
         return df.to_csv(index=False)
-    
+        
+    st.table(output_df)
     csv = convert_df(output_df)
     
     st.download_button(
@@ -110,5 +111,20 @@ else:
     try:
         output_df = pd.DataFrame(np.array(output).reshape(1, -1), columns=['input'] + [f'{i}th' for i in range(CFG.num_beams)] + ['valid compound'] + [f'{i}th score' for i in range(CFG.num_beams)] + ['valid compound score'])
         st.table(output_df)
+
+        @st.cache
+        def convert_df(df):
+            # IMPORTANT: Cache the conversion to prevent computation on every rerun
+            return df.to_csv(index=False)
+
+        csv = convert_df(output_df)
+        
+        st.download_button(
+            label="Download data as CSV",
+            data=csv,
+            file_name='output.csv',
+            mime='text/csv',
+        )
+        
     except:
         pass