Enhance model loading and error handling in app.py

app.py

@@ -126,7 +126,13 @@ def load_model_and_tokenizer():
                 status_text.text(f"⚠️ Failed to load from Hugging Face: {str(e)[:50]}...")
                 status_text.text("Loading base model as fallback...")
                 st.session_state.model = BigBirdForMaskedLM.from_pretrained("adibvafa/CodonTransformer")
-                st.session_state.model = st.session_state.model.to(st.session_state.device)
+                if isinstance(st.session_state.model, torch.nn.Module):
+                    if isinstance(st.session_state.model, torch.nn.Module):
+                        st.session_state.model.to(st.session_state.device)
+                    else:
+                        st.warning("Fallback model loaded is not a PyTorch module. Cannot move to device.")
+                else:
+                    st.warning("Fallback model loaded is not a PyTorch module. Cannot move to device.")
                 st.session_state.model_type = "base"
 
             progress_bar.progress(100)
@@ -296,7 +302,8 @@ def calculate_input_metrics(sequence: str, organism: str, sequence_type: str) ->
             metrics['tai'] = None
     try:
         analysis_dna = metrics['baseline_dna']
-        metrics['restriction_sites'] = len(scan_for_restriction_sites(analysis_dna))
+        # scan_for_restriction_sites returns an int, not a list, so no need for len()
+        metrics['restriction_sites'] = scan_for_restriction_sites(analysis_dna)
         metrics['negative_cis_elements'] = count_negative_cis_elements(analysis_dna)
         metrics['homopolymer_runs'] = calculate_homopolymer_runs(analysis_dna)
     except:
@@ -608,9 +615,9 @@ def run_optimization(protein: str, organism: str, use_post_processing: bool = Fa
                 # Create enhanced result object
                 from CodonTransformer.CodonUtils import DNASequencePrediction
                 st.session_state.post_processed_results = DNASequencePrediction(
-                    organism=result.organism,
-                    protein=result.protein,
-                    processed_input=result.processed_input,
+                    organism=_res.organism,
+                    protein=_res.protein,
+                    processed_input=_res.processed_input,
                     predicted_dna=polished_sequence
                 )
             except Exception as e:
@@ -706,7 +713,7 @@ def single_sequence_optimization():
         st.header("🧬 Input Sequence")
         sequence_input = st.text_area(
             "Enter Protein or DNA Sequence",
-            height=150,
+            height=300,
             placeholder="Enter protein sequence (MKWVT...) or DNA sequence (ATGGCG...)\n\nExample protein: MKWVTFISLLFLFSSAYSRGVFRRDAHKSEVAHRFKDLGEENFKALVLIAFAQYLQQCPFEDHVKLVNEVTEFAKTCVADESAENCDKSLHTLFGDKLCTVATLRETYGEMADCCAKQEPERNECFLQHKDDNPNLPRLVRPEVDVMCTAFHDNEETFLKKYLYEIARRHPYFYAPELLFFAKRYKAAFTECCQAADKAACLLPKLDELRDEGKASSAKQRLKCASLQKFGERAFKAWAVARLSQRFPKAEFAEVSKLVTDLTKVHTECCHGDLLECADDRADLAKYICENQDSISSKLKECCEKPLLEKSHCIAEVENDEMPADLPSLAADFVESKDVCKNYAEAKDVFLGMFLYEYARRHPDYSVVLLLRLAKTYETTLEKCCAAADPHECYAKVFDEFKPLVEEPQNLIKQNCELFEQLGEYKFQNALLVRYTKKVPQVSTPTLVEVSRNLGKVGSKCCKHPEAKRMPCAEDYLSVVLNQLCVLHEKTPVSDRVTKCCTE"
         )
         analyze_btn = st.button("Analyze Sequence", type="primary")
@@ -809,10 +816,10 @@ def single_sequence_optimization():
             if st.button("🚀 Optimize Sequence", type="primary", use_container_width=True):
                 st.session_state.results = None
                 if st.session_state.sequence_type == "dna":
-                    protein_sequence = translate_dna_to_protein(st.session_state.sequence_clean)
-                    run_optimization(protein_sequence, st.session_state.organism, use_post_processing)
+                    protein_sequence = translate_dna_to_protein(str(st.session_state.sequence_clean))
+                    run_optimization(protein_sequence, str(st.session_state.organism), use_post_processing)
                 else:
-                    run_optimization(st.session_state.sequence_clean, st.session_state.organism, use_post_processing)
+                    run_optimization(str(st.session_state.sequence_clean), str(st.session_state.organism), use_post_processing)
 
         # Enhanced progress display
         if st.session_state.optimization_running:
@@ -906,7 +913,7 @@ def display_optimization_results(result, organism, original_sequence, sequence_t
 
     # Optimized DNA sequence display
     st.subheader("🧬 Optimized DNA Sequence")
-    st.text_area("Optimized DNA Sequence", result.predicted_dna, height=100)
+    st.text_area("Optimized DNA Sequence", result.predicted_dna, height=300)
 
     # Enhanced download and export options
     col1, col2, col3 = st.columns(3)