Spaces:
Running
Running
File size: 3,303 Bytes
dad3b3d 13879a6 dad3b3d 13879a6 acbd0f9 13879a6 acbd0f9 dad3b3d 226fd5f dad3b3d 13879a6 acbd0f9 226fd5f acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 226fd5f acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 acbd0f9 13879a6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import streamlit as st
from merged import load_model_and_tokenizer, generate_code_stream
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():
return load_model_and_tokenizer()
tokenizer, model = load_model()
# ---------------- 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 ----------------
def select_language(lang):
st.session_state.selected_language = lang
st.session_state.generated_code = ""
st.session_state.pdf_bytes = b""
st.session_state.user_input = ""
def reset():
for k in st.session_state:
st.session_state[k] = "" if k != "selected_language" else None
# ---------------- UI ----------------
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.session_state.user_input = st.text_input(
"Describe the task",
value=st.session_state.user_input
)
# -------- 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()
# 1️⃣ Code Generation Loader
with st.spinner("⚙️ Generating code..."):
for token in generate_code_stream(
lang,
st.session_state.user_input,
tokenizer,
model
):
st.session_state.generated_code += token
code_placeholder.code(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()
# -------- 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"
)
|