Update src/streamlit_app.py

src/streamlit_app.py

@@ -308,7 +308,8 @@ elif input_method == "Paste Content":
         except Exception as e:
             st.sidebar.error(f"Error parsing content: {str(e)}")
 
-
+atoms.info["charge"] = 0
+atoms.info["spin"] = 0
 # Model selection
 st.sidebar.markdown("## Model Selection")
 model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem"])
@@ -326,6 +327,12 @@ if model_type == "FairChem":
     if selected_model == "UMA Small":
         st.sidebar.warning("Meta FAIR Acceptable Use Policy. This model was developed by the Fundamental AI Research (FAIR) team at Meta. By using it, you agree to their acceptable use policy, which prohibits using their models to violate the law or others' rights, plan or develop activities that present a risk of death or harm, and deceive or mislead others.")
         selected_task_type = st.sidebar.selectbox("Select UMA Model Task Type:", ["omol", "omat", "omc", "odac", "oc20"])
+        if selected_task_type == "omol":
+            charge = st.sidebar.number_input("Total Charge", min_value=-10, max_value=10, value=0)
+            spin_multiplicity = st.sidebar.number_input("Spin Multiplicity (2S + 1)", min_value=1, max_value=11)
+            # Set the total charge and spin multiplicity if using the OMol task
+            atoms.info["charge"] = charge
+            atoms.info["spin"] = spin_multiplicity
 # Check atom count limit
 if atoms is not None:
     check_atom_limit(atoms, selected_model)