Spaces:
Running
Running
| from __future__ import annotations | |
| import tempfile | |
| from pathlib import Path | |
| import streamlit as st | |
| from src.flashcard_generator.exporters import flashcards_to_apkg, flashcards_to_csv | |
| from src.flashcard_generator.extraction import extract_pdf_pages, extract_text | |
| from src.flashcard_generator.models import GenerationSettings | |
| from src.flashcard_generator.page_flashcards import format_page_flashcards, generate_page_flashcards | |
| from src.flashcard_generator.pipeline import FlashcardPipeline | |
| from src.flashcard_generator.text_processing import token_count | |
| st.set_page_config(page_title="AI Flashcard Generator", layout="wide") | |
| def main() -> None: | |
| st.title("AI-Powered Flashcard Generator") | |
| st.caption("Upload lecture notes, generate revision cards, and export an Anki-compatible deck.") | |
| with st.sidebar: | |
| st.header("Generation") | |
| model_name = st.selectbox( | |
| "Hugging Face model", | |
| ["google/flan-t5-small", "facebook/bart-large-cnn", "t5-small"], | |
| index=0, | |
| ) | |
| use_model = st.toggle("Use Hugging Face summariser", value=True) | |
| generate_all_possible = st.toggle("Generate all possible Q/A", value=False) | |
| if generate_all_possible: | |
| st.caption("Uses every detected concept and every question style.") | |
| concepts_per_chunk = 0 | |
| cards_per_concept = 10 | |
| else: | |
| concepts_per_chunk = st.slider("Concepts per chunk", 3, 20, 8) | |
| cards_per_concept = st.slider("Questions per concept", 1, 10, 3) | |
| harden_questions = st.toggle("Rewrite with Bloom's Taxonomy", value=True) | |
| uploaded = st.file_uploader( | |
| "Upload a PDF, text file, or lecture-note image", | |
| type=["pdf", "txt", "md", "png", "jpg", "jpeg", "webp", "tiff", "bmp"], | |
| ) | |
| if uploaded is None: | |
| st.info("Choose a lecture file to begin.") | |
| return | |
| suffix = Path(uploaded.name).suffix | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temporary: | |
| temporary.write(uploaded.getbuffer()) | |
| temp_path = Path(temporary.name) | |
| is_pdf = suffix.lower() == ".pdf" | |
| page_texts: list[str] = [] | |
| try: | |
| with st.spinner("Extracting text..."): | |
| if is_pdf: | |
| page_texts = extract_pdf_pages(temp_path) | |
| text = "\n".join(page_texts) | |
| else: | |
| text = extract_text(temp_path) | |
| except Exception as exc: | |
| st.error(f"Could not extract text: {exc}") | |
| return | |
| finally: | |
| temp_path.unlink(missing_ok=True) | |
| if not text: | |
| st.warning("No readable text was found in the uploaded file.") | |
| return | |
| st.subheader("Extracted Notes") | |
| metric_left, metric_right = st.columns(2) | |
| metric_left.metric("Words", token_count(text)) | |
| metric_right.metric("Pages" if is_pdf else "Characters", len(page_texts) if is_pdf else len(text)) | |
| with st.expander("Preview extracted text", expanded=False): | |
| st.write(text[:5000]) | |
| if not st.button("Generate Flashcards", type="primary"): | |
| return | |
| if is_pdf: | |
| with st.spinner("Building page-by-page flashcards..."): | |
| page_sets = generate_page_flashcards(page_texts, questions_per_page=10) | |
| formatted_output = format_page_flashcards(page_sets) | |
| cards = [card for page_set in page_sets for card in page_set.cards] | |
| if not cards: | |
| st.warning("No flashcards could be generated from this PDF.") | |
| return | |
| st.success(f"Generated {len(cards)} flashcards from {len(page_sets)} page(s).") | |
| tab_output, tab_cards, tab_export = st.tabs(["Formatted Output", "Flashcards", "Export"]) | |
| with tab_output: | |
| st.text(formatted_output) | |
| with tab_cards: | |
| for page_set in page_sets: | |
| st.markdown(f"**Page {page_set.page_number}**") | |
| for index, card in enumerate(page_set.cards, start=1): | |
| with st.container(border=True): | |
| st.markdown(f"**Q{index}: {card.question}**") | |
| st.write(f"A{index}: {card.short_answer}") | |
| if page_set.insufficient_note: | |
| st.caption(page_set.insufficient_note) | |
| with tab_export: | |
| st.download_button( | |
| "Download TXT", | |
| data=formatted_output, | |
| file_name="page_flashcards.txt", | |
| mime="text/plain", | |
| ) | |
| csv_data = flashcards_to_csv(cards) | |
| st.download_button( | |
| "Download CSV for Anki", | |
| data=csv_data, | |
| file_name="flashcards.csv", | |
| mime="text/csv", | |
| ) | |
| try: | |
| apkg_data = flashcards_to_apkg(cards) | |
| st.download_button( | |
| "Download APKG", | |
| data=apkg_data, | |
| file_name="flashcards.apkg", | |
| mime="application/octet-stream", | |
| ) | |
| except Exception as exc: | |
| st.caption(f"APKG export unavailable: {exc}") | |
| return | |
| settings = GenerationSettings( | |
| concepts_per_chunk=concepts_per_chunk, | |
| cards_per_concept=cards_per_concept, | |
| generate_all_possible=generate_all_possible, | |
| harden_questions=harden_questions, | |
| model_name=model_name, | |
| use_model=use_model, | |
| ) | |
| with st.spinner("Summarising notes and building flashcards..."): | |
| pipeline = FlashcardPipeline(settings) | |
| cards, summaries, concepts = pipeline.run(text) | |
| if not cards: | |
| st.warning("No flashcards could be generated from this document.") | |
| return | |
| st.success(f"Generated {len(cards)} flashcards from {len(summaries)} chunk(s).") | |
| tab_cards, tab_summaries, tab_concepts, tab_export = st.tabs( | |
| ["Flashcards", "Summaries", "Concepts JSON", "Export"] | |
| ) | |
| with tab_cards: | |
| for index, card in enumerate(cards, start=1): | |
| with st.container(border=True): | |
| st.markdown(f"**{index}. {card.question}**") | |
| st.markdown("**Short answer**") | |
| st.write(card.short_answer) | |
| with st.expander("Long answer", expanded=False): | |
| st.write(card.long_answer) | |
| st.caption(f"{card.concept} - {card.difficulty} - {card.bloom_level}") | |
| with tab_summaries: | |
| for index, summary in enumerate(summaries, start=1): | |
| st.markdown(f"**Chunk {index}**") | |
| st.write(summary) | |
| with tab_concepts: | |
| st.json(concepts) | |
| with tab_export: | |
| csv_data = flashcards_to_csv(cards) | |
| st.download_button( | |
| "Download CSV for Anki", | |
| data=csv_data, | |
| file_name="flashcards.csv", | |
| mime="text/csv", | |
| ) | |
| try: | |
| apkg_data = flashcards_to_apkg(cards) | |
| st.download_button( | |
| "Download APKG", | |
| data=apkg_data, | |
| file_name="flashcards.apkg", | |
| mime="application/octet-stream", | |
| ) | |
| except Exception as exc: | |
| st.caption(f"APKG export unavailable: {exc}") | |
| if __name__ == "__main__": | |
| main() | |