FINAL DOCKERFIX: Ensuring Dockerfile copies pdf_utils.py instead of the nonexistent main.py.

app.py

@@ -1,8 +1,11 @@
 import streamlit as st
 from merged import load_model_and_tokenizer, generate_code_stream
-from pdf_utils import test_case   
+from pdf_utils import test_case
 
+# ---------------- App Boot ----------------
+st.set_page_config(page_title="Code + Test Case Generator", layout="wide")
 st.write("🚀 App started")
+
 # ---------------- Load model ONCE ----------------
 @st.cache_resource
 def load_model():
@@ -10,16 +13,16 @@ def load_model():
 
 tokenizer, model = load_model()
 
-# ---------------- Session state ----------------
+# ---------------- Session State ----------------
 for key in ["selected_language", "generated_code", "pdf_bytes", "user_input"]:
     if key not in st.session_state:
         st.session_state[key] = "" if key != "selected_language" else None
 
-# ---------------- UI helpers ----------------
+# ---------------- UI Helpers ----------------
 def select_language(lang):
     st.session_state.selected_language = lang
     st.session_state.generated_code = ""
-    st.session_state.pdf_bytes = ""
+    st.session_state.pdf_bytes = b""
     st.session_state.user_input = ""
 
 def reset():
@@ -27,33 +30,40 @@ def reset():
         st.session_state[k] = "" if k != "selected_language" else None
 
 # ---------------- UI ----------------
-st.title("Generate any code and get test case for it")
+st.title("🧠 Generate Code & Download Test Cases")
 
+# -------- Language Selection --------
 if st.session_state.selected_language is None:
+    st.subheader("Select Language")
     cols = st.columns(4)
     cols[0].button("Java", on_click=select_language, args=("Java",))
     cols[1].button("React.js", on_click=select_language, args=("React.js",))
     cols[2].button("Python", on_click=select_language, args=("Python",))
     cols[3].button("C++", on_click=select_language, args=("C++",))
 
+# -------- Main App --------
 else:
     lang = st.session_state.selected_language
     st.subheader(f"Selected Language: {lang}")
-    st.button("Reset", on_click=reset)
+    st.button("🔄 Reset", on_click=reset)
 
     st.session_state.user_input = st.text_input(
         "Describe the task",
         value=st.session_state.user_input
     )
 
-    if st.button("Generate Code"):
+    # -------- Generate Code Button --------
+    if st.button("⚙️ Generate Code"):
         if st.session_state.user_input.strip():
+
             st.session_state.generated_code = ""
+            st.session_state.pdf_bytes = b""
 
             st.subheader("Generated Code")
             code_placeholder = st.empty()
 
-            with st.spinner("Generating code..."):
+            # 1️⃣ Code Generation Loader
+            with st.spinner("⚙️ Generating code..."):
                 for token in generate_code_stream(
                     lang,
                     st.session_state.user_input,
@@ -61,27 +71,27 @@ else:
                     model
                 ):
                     st.session_state.generated_code += token
-                    code_placeholder.code(
-                        st.session_state.generated_code
-                    )
+                    code_placeholder.code(st.session_state.generated_code)
 
-            # Generate test case PDF AFTER code generation
-            st.session_state.pdf_bytes = test_case(
-                st.session_state.generated_code
-            )
+            # 2️⃣ PDF Generation Loader
+            with st.spinner("📄 Generating test cases PDF..."):
+                st.session_state.pdf_bytes = test_case(
+                    st.session_state.generated_code
+                )
 
+            st.success("✅ Test cases PDF generated")
 
         else:
             st.warning("Please enter a task")
 
     st.divider()
 
-    with st.spinner("📄 Generating test cases PDF..."):
-        st.session_state.pdf_bytes = test_case(
-        st.session_state.generated_code
-    )
-
-    st.success("✅ Test cases PDF generated")
-
-
-
+    # -------- Download PDF --------
+    if st.session_state.get("pdf_bytes"):
+        st.download_button(
+            label="📥 Download Test Cases PDF",
+            data=st.session_state.pdf_bytes,
+            file_name="test_cases.pdf",
+            mime="application/pdf",
+            key="download_pdf"
+        )