Update app.py

app.py

@@ -1,98 +1,79 @@
 import streamlit as st
+import os
+import tempfile
+import fitz
 from config import (
-    process_pdf, process_image, process_audio, process_text, process_youtube,
-    generate_flashcards, export_to_apkg, export_to_csv
+    process_pdf,
+    process_text,
+    process_image,
+    process_audio,
+    process_youtube,
+    generate_flashcards,
+    export_to_apkg,
+    export_to_csv,
 )
-import tempfile
-import os
-import uuid
-
-st.set_page_config(page_title="🦇 BatAnki Pro", layout="wide")
-st.title("🦇 BatAnki – AI Flashcard Engine for Learners")
 
-st.sidebar.title("📁 Upload Source")
-upload_type = st.sidebar.radio("Select input type", ["PDF", "Image", "Audio", "Text File", "YouTube Link"])
-deck_name = st.sidebar.text_input("📌 Deck name", value="BatAnki_Deck")
-model_choice = st.sidebar.selectbox("🧠 AI Model", ["t5-base", "microsoft/BioGPT"])
-generate_types = st.sidebar.multiselect("Card Types", ["Q&A", "Cloze", "MCQ", "Reverse"], default=["Q&A"])
+st.set_page_config(page_title="🦇 BatAnki AI", layout="wide")
+st.title("📚 BatAnki – AI Flashcard Generator")
 
-uploaded_file = None
-yt_url = ""
+st.sidebar.markdown("📁 **Input Options**")
+uploaded_file = st.sidebar.file_uploader("Upload a file", type=["pdf", "txt", "docx", "jpg", "png", "mp3", "wav"])
+youtube_url = st.sidebar.text_input("Or paste YouTube URL")
 
-if upload_type == "YouTube Link":
-    yt_url = st.sidebar.text_input("📺 YouTube URL")
-else:
-    uploaded_file = st.sidebar.file_uploader("Upload file", type=["pdf", "png", "jpg", "jpeg", "mp3", "wav", "txt", "docx"])
+deck_name = st.text_input("Deck Name", value="BatAnkiDeck")
+types_selected = st.multiselect("Card Types", ["Q&A", "Cloze", "MCQ", "Reverse"], default=["Q&A"])
+max_cards = st.slider("Max Cards", 5, 1000, 50)
 
-# Extract text from the file
-extracted_text = ""
-file_path = None
+input_text = ""
+cards = []
 
 if uploaded_file:
-    with tempfile.NamedTemporaryFile(delete=False) as tmp:
-        tmp.write(uploaded_file.read())
-        file_path = tmp.name
+    suffix = uploaded_file.name.split(".")[-1]
+    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{suffix}") as tmp_file:
+        tmp_file.write(uploaded_file.read())
+        tmp_path = tmp_file.name
 
-if st.sidebar.button("🧠 Extract Text"):
-    if upload_type == "PDF":
-        extracted_text = process_pdf(file_path)
-    elif upload_type == "Image":
-        extracted_text = process_image(file_path)
-    elif upload_type == "Audio":
-        extracted_text = process_audio(file_path)
-    elif upload_type == "Text File":
-        extracted_text = process_text(file_path)
-    elif upload_type == "YouTube Link" and yt_url:
-        extracted_text = process_youtube(yt_url)
-    else:
-        st.warning("Please upload or enter a valid input.")
+    if suffix == "pdf":
+        doc = fitz.open(tmp_path)
+        st.info("PDF preview:")
+        page_number = st.number_input("Page", 1, len(doc), 1)
+        page = doc[page_number - 1]
+        pix = page.get_pixmap()
+        st.image(pix.tobytes("png"), caption=f"Page {page_number}")
+        text = page.get_text()
+        input_text = text if text.strip() else process_pdf(tmp_path)
+        if st.button("Generate Cards from This Page"):
+            cards = generate_flashcards(input_text, types_selected, max_cards)
+    elif suffix in ["jpg", "png"]:
+        input_text = process_image(tmp_path)
+        st.image(tmp_path)
+    elif suffix in ["mp3", "wav"]:
+        input_text = process_audio(tmp_path)
+    elif suffix in ["txt", "docx"]:
+        input_text = process_text(tmp_path)
 
-    if not extracted_text.strip():
-        st.error("No text extracted.")
-    else:
-        st.success("✅ Text extracted!")
-        st.session_state["extracted_text"] = extracted_text
-        st.session_state["deck_name"] = deck_name
-        st.session_state["generate_types"] = generate_types
-        st.session_state["model"] = model_choice
+elif youtube_url:
+    st.info("Processing YouTube audio...")
+    input_text = process_youtube(youtube_url)
 
-if "extracted_text" in st.session_state:
-    st.subheader("📄 Extracted Text")
-    st.text_area("Extracted Content", st.session_state["extracted_text"][:10000], height=300)
+if input_text and not cards:
+    if st.button("Generate Cards"):
+        cards = generate_flashcards(input_text, types_selected, max_cards)
 
-    if st.button("🚀 Generate Flashcards"):
-        with st.spinner("Generating cards..."):
-            cards = generate_flashcards(
-                st.session_state["extracted_text"],
-                model_name=st.session_state["model"],
-                types=st.session_state["generate_types"]
-            )
-            st.session_state["cards"] = cards
-        st.success(f"✅ Generated {len(cards)} cards")
-
-if "cards" in st.session_state:
-    st.subheader("🧾 Flashcard Preview & Tagging")
-    tagged_cards = []
-    for idx, card in enumerate(st.session_state["cards"], 1):
-        st.markdown(f"**Card {idx}**")
-        question = st.text_input(f"Q{idx}", card["question"])
-        answer = st.text_area(f"A{idx}", card["answer"])
-        tag = st.text_input(f"🏷️ Tag for Card {idx}", value=card.get("tag", ""))
-        tagged_cards.append({"question": question, "answer": answer, "tag": tag})
+if cards:
+    st.subheader("Generated Flashcards")
+    for i, card in enumerate(cards):
+        st.markdown(f"**{i+1}. {card['question']}**")
+        st.markdown(f"*Answer:* {card['answer']}")
         st.markdown("---")
-    st.session_state["cards"] = tagged_cards
 
-    st.subheader("📦 Export")
     col1, col2 = st.columns(2)
-
     with col1:
-        if st.button("📥 Export to Anki (.apkg)"):
-            path = export_to_apkg(st.session_state["cards"], st.session_state["deck_name"])
-            with open(path, "rb") as f:
-                st.download_button("⬇️ Download APKG", f, f"{st.session_state['deck_name']}.apkg")
-
+        if st.button("Export to CSV"):
+            export_to_csv(cards)
+            st.success("Exported CSV")
     with col2:
-        if st.button("📊 Export to CSV"):
-            path = export_to_csv(st.session_state["cards"], st.session_state["deck_name"])
+        if st.button("Export to Anki (.apkg)"):
+            path = export_to_apkg(cards, deck_name)
             with open(path, "rb") as f:
-                st.download_button("⬇️ Download CSV", f, f"{st.session_state['deck_name']}.csv")
+                st.download_button("Download .apkg", f, file_name=path)