refactor

app.py

@@ -108,43 +108,40 @@ with col2:
         else:
             try:
                 selected_model = models[model_choice]
-                
-                if selected_model is None:
-                    st.error("Selected model is not implemented yet")
-                else:
-                    with st.spinner("Generating summary..."):
-                        if model_choice == "mt5-small":
-                            result = selected_model(
-                                email_input,
-                                max_length=150,
-                                do_sample=True,
-                                repetition_penalty=1.5
-                            )[0]['summary_text']
-                        elif model_choice == "Llama 3.2":
-                            model_obj, tokenizer = selected_model
-                            result = generate_llama_summary(
-                                email_input,
-                                model_obj,
-                                tokenizer,
-                                PROMPT_TEMPLATE
-                            )
-                        
-                        # Display results
-                        st.success("**Generated Summary:**")
-                        st.write(result)
-                        
-                        # Show reference summary if example is selected
-                        if example_choice in EXAMPLES:
-                            st.success("**Reference Summary:**")
-                            st.write(EXAMPLES[example_choice]["reference"])
-                        
-                        # Add export options
-                        st.download_button(
-                            label="Download Summary",
-                            data=result,
-                            file_name="email_summary.txt",
-                            mime="text/plain"
+
+                with st.spinner("Generating summary..."):
+                    if model_choice == "mt5-small":
+                        result = selected_model(
+                            email_input,
+                            max_length=150,
+                            do_sample=True,
+                            repetition_penalty=1.5
+                        )[0]['summary_text']
+                    elif model_choice == "Llama 3.2":
+                        model_obj, tokenizer = selected_model
+                        result = generate_llama_summary(
+                            email_input,
+                            model_obj,
+                            tokenizer,
+                            PROMPT_TEMPLATE
                         )
+                    
+                    # Display results
+                    st.success("**Generated Summary:**")
+                    st.write(result)
+                    
+                    # Show reference summary if example is selected
+                    if example_choice in EXAMPLES:
+                        st.success("**Reference Summary:**")
+                        st.write(EXAMPLES[example_choice]["reference"])
+                    
+                    # Add export options
+                    st.download_button(
+                        label="Download Summary",
+                        data=result,
+                        file_name="email_summary.txt",
+                        mime="text/plain"
+                    )
                         
             except Exception as e:
                 st.error(f"Error generating summary: {str(e)}")