Spaces:
Build error
Build error
| import os | |
| import re | |
| import warnings | |
| import io | |
| from pathlib import Path | |
| from datetime import datetime | |
| os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' | |
| os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0' | |
| warnings.filterwarnings("ignore", category=DeprecationWarning) | |
| warnings.filterwarnings("ignore", category=UserWarning) | |
| import streamlit as st | |
| import chromadb | |
| from sentence_transformers import SentenceTransformer | |
| from huggingface_hub import snapshot_download, InferenceClient | |
| from reportlab.lib.pagesizes import A4 | |
| from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle | |
| from reportlab.lib.units import cm | |
| from reportlab.lib import colors | |
| from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, HRFlowable, Image as RLImage | |
| from reportlab.lib.enums import TA_LEFT, TA_CENTER | |
| # ── CONFIGURE ──────────────────────────────────────────────── | |
| HF_DB_REPO = "Daksh17440/remote_sensing_db" | |
| DB_LOCAL_PATH = "/tmp/rs_data" | |
| DB_PATH = "/tmp/rs_data/vector_db_complete" | |
| MINERU_PATH = "/tmp/rs_data/mineru_output_complete" | |
| COLLECTION_NAME = "rs_index" | |
| HF_LLM_MODEL = "mistralai/Mistral-7B-Instruct-v0.3" | |
| EMBED_MODEL = "all-MiniLM-L6-v2" | |
| TOP_K = 5 | |
| L2_THRESHOLD = 1.2 | |
| # ───────────────────────────────────────────────────────────── | |
| CHAPTER_INFO = { | |
| "RAS_ch1": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 1 — The Nature of Remote Sensing"), | |
| "RAS_ch2": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 2 — Optical Radiation Models"), | |
| "RAS_ch3": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 3 — Sensor Models"), | |
| "RAS_ch4": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 4 — Data Models"), | |
| "RAS_ch5": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 5 — Spectral Transforms"), | |
| "RAS_ch6": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 6 — Spatial Transforms"), | |
| "RAS_ch7": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 7 — Correction and Calibration"), | |
| "RAS_ch8": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 8 — Image Registration and Fusion"), | |
| "RAS_ch9": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Chapter 9 — Thematic Classification"), | |
| "RAS_fig": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Figures"), | |
| "RAS_tables": ("Remote Sensing: Models and Methods for Image Processing (Schowengerdt)", "Tables"), | |
| "RRS_ch1": ("Radar Remote Sensing (2022)", "Chapter 1 — Introduction to RADAR Remote Sensing"), | |
| "RRS_ch2": ("Radar Remote Sensing (2022)", "Chapter 2 — Microwave Components and Devices"), | |
| "RRS_ch3": ("Radar Remote Sensing (2022)", "Chapter 3 — Theory of Monostatic and Bistatic Radar"), | |
| "RRS_ch4": ("Radar Remote Sensing (2022)", "Chapter 4 — Review of Microwave Fundamentals"), | |
| "RRS_ch5": ("Radar Remote Sensing (2022)", "Chapter 5 — Comparative Flood Area Analysis"), | |
| "RRS_ch6": ("Radar Remote Sensing (2022)", "Chapter 6 — Subsurface Feature Identification Using L-Band SAR"), | |
| "RRS_ch7": ("Radar Remote Sensing (2022)", "Chapter 7 — Terrestrial Water Budget Through Radar Remote Sensing"), | |
| "RRS_ch8": ("Radar Remote Sensing (2022)", "Chapter 8 — Application of SAR Remote Sensing"), | |
| "RRS_ch9": ("Radar Remote Sensing (2022)", "Chapter 9 — Classification of Radar Data Using Bayesian Optimization"), | |
| "RRS_ch10": ("Radar Remote Sensing (2022)", "Chapter 10 — Modeling and Simulation of SAR"), | |
| "RRS_ch11": ("Radar Remote Sensing (2022)", "Chapter 11 — Flood Inundation Mapping from SAR"), | |
| "RRS_ch12": ("Radar Remote Sensing (2022)", "Chapter 12 — Performance Assessment of Phased Array L-Band SAR"), | |
| "RRS_ch13": ("Radar Remote Sensing (2022)", "Chapter 13 — Evaluation of Speckle Filtering Methods"), | |
| "RRS_ch14": ("Radar Remote Sensing (2022)", "Chapter 14 — Emerging Techniques of Polarimetric Interferometry"), | |
| "RRS_ch15": ("Radar Remote Sensing (2022)", "Chapter 15 — Advanced Method for Radar Remote Sensing"), | |
| "RRS_ch16": ("Radar Remote Sensing (2022)", "Chapter 16 — Estimating Crop Biophysical Parameters"), | |
| "RRS_ch17": ("Radar Remote Sensing (2022)", "Chapter 17 — Fuzzy Logic for Retrieval of Kidney Bean Crop"), | |
| "RRS_ch18": ("Radar Remote Sensing (2022)", "Chapter 18 — Monitoring Tropical Peatlands Subsidence"), | |
| "RRS_ch19": ("Radar Remote Sensing (2022)", "Chapter 19 — North American Continental Wetland Map"), | |
| "RRS_ch20": ("Radar Remote Sensing (2022)", "Chapter 20 — Challenges in Radar Remote Sensing"), | |
| "RRS_ch21": ("Radar Remote Sensing (2022)", "Chapter 21 — ISRO SAR Satellite Study"), | |
| "RRS_ch22": ("Radar Remote Sensing (2022)", "Chapter 22 — Radar Remote Sensing of Soil Moisture"), | |
| "ORS_ch1": ("Optical Remote Sensing of Land Surface (2017)", "Chapter 1 — Radiometry in the Optical Domain"), | |
| "ORS_ch2": ("Optical Remote Sensing of Land Surface (2017)", "Chapter 2 — Multispectral Satellite Image Processing"), | |
| "ORS_ch3": ("Optical Remote Sensing of Land Surface (2017)", "Chapter 3 — Digital Terrain Models from Optical Data"), | |
| "ORS_ch4": ("Optical Remote Sensing of Land Surface (2017)", "Chapter 4 — Processing Hyperspectral Images"), | |
| "ORS_ch5": ("Optical Remote Sensing of Land Surface (2017)", "Chapter 5 — Principle and Physics of the LiDAR Measurement"), | |
| "ORS_ch6": ("Optical Remote Sensing of Land Surface (2017)", "Chapter 6 — Airborne LiDAR Data Processing"), | |
| "ORS_ch7": ("Optical Remote Sensing of Land Surface (2017)", "Chapter 7 — Digital Terrain Models Derived from Airborne LiDAR"), | |
| "LSR_ch1": ("Land Surface Remote Sensing in Urban and Coastal Areas (2017)", "Chapter 1 — Optical Remote Sensing in Urban Areas"), | |
| "LSR_ch2": ("Land Surface Remote Sensing in Urban and Coastal Areas (2017)", "Chapter 2 — Urban Scene Analysis with Mobile Mapping"), | |
| "LSR_ch3": ("Land Surface Remote Sensing in Urban and Coastal Areas (2017)", "Chapter 3 — Satellite Imagery as a Tool for Terrain Analysis"), | |
| "LSR_ch4": ("Land Surface Remote Sensing in Urban and Coastal Areas (2017)", "Chapter 4 — Remote Sensing and Ocean"), | |
| "LSR_ch5": ("Land Surface Remote Sensing in Urban and Coastal Areas (2017)", "Chapter 5 — LiDAR Measurements and Applications"), | |
| "LSR_ch6": ("Land Surface Remote Sensing in Urban and Coastal Areas (2017)", "Chapter 6 — Contributions of Airborne Topographic LiDAR"), | |
| "LSR_ch7": ("Land Surface Remote Sensing in Urban and Coastal Areas (2017)", "Chapter 7 — Mangrove Forest Dynamics Using Very High Resolution"), | |
| "LSR_ch8": ("Land Surface Remote Sensing in Urban and Coastal Areas (2017)", "Chapter 8 — Remote Sensing Based Monitoring"), | |
| "ARS_ch1": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 1 — A Systematic View of Remote Sensing"), | |
| "ARS_ch2": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 2 — Geometric Processing and Positioning Techniques"), | |
| "ARS_ch3": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 3 — Compositing, Smoothing and Gap Filling"), | |
| "ARS_ch4": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 4 — Data Fusion"), | |
| "ARS_ch5": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 5 — Atmospheric Correction of Optical Images"), | |
| "ARS_ch6": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 6 — Solar Radiation"), | |
| "ARS_ch7": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 7 — Broadband Albedo"), | |
| "ARS_ch8": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 8 — Land Surface Temperature and Thermal Infrared"), | |
| "ARS_ch9": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 9 — Surface Longwave Radiation Budget"), | |
| "ARS_ch10": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 10 — Canopy Biochemical Characteristics"), | |
| "ARS_ch11": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 11 — Leaf Area Index"), | |
| "ARS_ch12": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 12 — Fraction of Absorbed Photosynthetically Active Radiation"), | |
| "ARS_ch13": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 13 — Fractional Vegetation Cover"), | |
| "ARS_ch14": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 14 — Vegetation Height and Vertical Structure"), | |
| "ARS_ch15": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 15 — Above-ground Biomass"), | |
| "ARS_ch16": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 16 — Vegetation Production in Terrestrial Ecosystems"), | |
| "ARS_ch17": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 17 — Precipitation"), | |
| "ARS_ch18": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 18 — Terrestrial Evapotranspiration"), | |
| "ARS_ch19": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 19 — Soil Moisture Content"), | |
| "ARS_ch20": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 20 — Snow Water Equivalence"), | |
| "ARS_ch21": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 21 — Water Storage"), | |
| "ARS_ch22": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 22 — High-level Land Product Integration"), | |
| "ARS_ch23": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 23 — Production and Data Management Systems"), | |
| "ARS_ch24": ("Advanced Remote Sensing (Liang, 2012)", "Chapter 24 — Land Cover and Land Use Changes"), | |
| "ERS_ch1": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 1 — Introduction"), | |
| "ERS_ch2": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 2 — Balloon-borne Radiometers"), | |
| "ERS_ch3": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 3 — Frost Point Hygrometers"), | |
| "ERS_ch4": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 4 — Ozonesondes"), | |
| "ERS_ch5": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 5 — Oceanographic Buoys"), | |
| "ERS_ch6": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 6 — Surface-based Thermal Infrared"), | |
| "ERS_ch7": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 7 — Sun Photometers"), | |
| "ERS_ch8": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 8 — The AirCore Atmospheric Sampler"), | |
| "ERS_ch9": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 9 — High Altitude Aircraft Radiometers"), | |
| "ERS_ch10": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 10 — Aircraft Dropsondes"), | |
| "ERS_ch11": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 11 — Ship-based Cal/Val"), | |
| "ERS_ch12": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 12 — Land-based Cal/Val"), | |
| "ERS_ch13": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 13 — Aircraft Vertical Profile Measurements"), | |
| "ERS_ch14": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 14 — Campaign Situational Awareness"), | |
| "ERS_ch15": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 15 — On-orbit VIIRS Sensor Calibration"), | |
| "ERS_ch16": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 16 — The NOAA Sounding Products"), | |
| "ERS_ch17": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 17 — Satellite Microwave Sounding"), | |
| "ERS_ch18": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 18 — Considerations for Thermal Sensors"), | |
| "ERS_ch19": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 19 — Sea Surface Temperature Validation"), | |
| "ERS_ch20": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 20 — Satellite Ocean Color"), | |
| "ERS_ch21": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 21 — Land Surface Temperature"), | |
| "ERS_ch22": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 22 — Heterogeneity of Smoke from Fires"), | |
| "ERS_ch23": ("Field Measurements for Passive Environmental Remote Sensing (2022)", "Chapter 23 — Downburst Monitoring"), | |
| "CL25_4231": ("NASA-ISRO SAR (NISAR) Mission Science Users Handbook", "NASA"), | |
| "RADAR_INTERFEROMETRY": ("RADAR INTERFEROMETRY", "Remote Sensing and Digital Image Processing") | |
| } | |
| def get_citation(source_file, chapter_name_meta): | |
| key = chapter_name_meta or Path(source_file).stem | |
| if key in CHAPTER_INFO: | |
| return CHAPTER_INFO[key] | |
| stem = Path(source_file).stem | |
| if stem in CHAPTER_INFO: | |
| return CHAPTER_INFO[stem] | |
| return (source_file.replace('.pdf','').replace('_',' '), "") | |
| def find_image_in_mineru(image_path_meta, chapter_name): | |
| """ | |
| Find actual image file from mineru_output_complete. | |
| image_path_meta: path stored in ChromaDB metadata (may be old local path) | |
| chapter_name: e.g. RAS_ch1 | |
| """ | |
| if not image_path_meta: | |
| return None | |
| # First try: use stored path directly (works if DB was built on same machine) | |
| p = Path(image_path_meta) | |
| if p.exists(): | |
| return str(p) | |
| # Second try: extract just the filename and find it in mineru output | |
| img_filename = Path(image_path_meta).name | |
| if not img_filename: | |
| return None | |
| # Search in the chapter's mineru folder | |
| chapter_dir = Path(MINERU_PATH) / chapter_name | |
| if chapter_dir.exists(): | |
| matches = list(chapter_dir.rglob(img_filename)) | |
| if matches: | |
| return str(matches[0]) | |
| # Third try: search all mineru output by filename | |
| all_matches = list(Path(MINERU_PATH).rglob(img_filename)) | |
| if all_matches: | |
| return str(all_matches[0]) | |
| return None | |
| def render_answer(text): | |
| parts = re.split(r'(\[.*?\]|\$\$.*?\$\$)', text, flags=re.DOTALL) | |
| for part in parts: | |
| part = part.strip() | |
| if not part: | |
| continue | |
| if (part.startswith('[') and part.endswith(']')) or \ | |
| (part.startswith('$$') and part.endswith('$$')): | |
| latex = part.strip('[]').strip('$$').strip() | |
| try: | |
| st.latex(latex) | |
| except Exception: | |
| st.markdown(part) | |
| else: | |
| st.markdown(part) | |
| def generate_chat_pdf(messages): | |
| """Generate a PDF of the full chat history using ReportLab.""" | |
| buffer = io.BytesIO() | |
| doc = SimpleDocTemplate( | |
| buffer, | |
| pagesize=A4, | |
| rightMargin=2*cm, leftMargin=2*cm, | |
| topMargin=2*cm, bottomMargin=2*cm | |
| ) | |
| styles = getSampleStyleSheet() | |
| title_style = ParagraphStyle( | |
| 'Title', parent=styles['Title'], | |
| fontSize=18, textColor=colors.HexColor('#1F3864'), | |
| spaceAfter=6 | |
| ) | |
| subtitle_style = ParagraphStyle( | |
| 'Subtitle', parent=styles['Normal'], | |
| fontSize=10, textColor=colors.grey, | |
| spaceAfter=20 | |
| ) | |
| user_label_style = ParagraphStyle( | |
| 'UserLabel', parent=styles['Normal'], | |
| fontSize=9, textColor=colors.white, | |
| backColor=colors.HexColor('#2E75B6'), | |
| leftIndent=0, borderPadding=4, | |
| spaceAfter=4, spaceBefore=12 | |
| ) | |
| user_text_style = ParagraphStyle( | |
| 'UserText', parent=styles['Normal'], | |
| fontSize=11, textColor=colors.HexColor('#1a1a1a'), | |
| leftIndent=10, spaceAfter=8, | |
| backColor=colors.HexColor('#EBF3FB'), | |
| borderPadding=6 | |
| ) | |
| assistant_label_style = ParagraphStyle( | |
| 'AssistantLabel', parent=styles['Normal'], | |
| fontSize=9, textColor=colors.white, | |
| backColor=colors.HexColor('#1F3864'), | |
| leftIndent=0, borderPadding=4, | |
| spaceAfter=4, spaceBefore=12 | |
| ) | |
| assistant_text_style = ParagraphStyle( | |
| 'AssistantText', parent=styles['Normal'], | |
| fontSize=11, textColor=colors.HexColor('#1a1a1a'), | |
| leftIndent=10, spaceAfter=8, | |
| borderPadding=6 | |
| ) | |
| citation_style = ParagraphStyle( | |
| 'Citation', parent=styles['Normal'], | |
| fontSize=9, textColor=colors.HexColor('#404040'), | |
| leftIndent=20, spaceAfter=3, | |
| italics=True | |
| ) | |
| citation_header_style = ParagraphStyle( | |
| 'CitationHeader', parent=styles['Normal'], | |
| fontSize=9, textColor=colors.HexColor('#2E75B6'), | |
| leftIndent=10, spaceAfter=3, | |
| fontName='Helvetica-Bold' | |
| ) | |
| story = [] | |
| # Title | |
| story.append(Paragraph("Remote Sensing RAG — Chat Export", title_style)) | |
| story.append(Paragraph( | |
| f"Exported on {datetime.now().strftime('%B %d, %Y at %H:%M')}", | |
| subtitle_style | |
| )) | |
| story.append(HRFlowable(width="100%", thickness=2, color=colors.HexColor('#2E75B6'))) | |
| story.append(Spacer(1, 0.4*cm)) | |
| for msg in messages: | |
| role = msg["role"] | |
| if role == "user": | |
| story.append(Paragraph("🧑 You", user_label_style)) | |
| # Clean text for PDF (strip markdown symbols) | |
| clean_text = msg["content"].replace('**','').replace('*','').replace('#','') | |
| story.append(Paragraph(clean_text, user_text_style)) | |
| elif role == "assistant": | |
| story.append(Paragraph("🛰️ Assistant", assistant_label_style)) | |
| clean_text = msg["content"].replace('**','').replace('*','').replace('#','') | |
| # Strip LaTeX blocks for PDF (replace with placeholder) | |
| clean_text = re.sub(r'\[.*?\]', '[equation]', clean_text, flags=re.DOTALL) | |
| clean_text = re.sub(r'\$\$.*?\$\$', '[equation]', clean_text, flags=re.DOTALL) | |
| story.append(Paragraph(clean_text, assistant_text_style)) | |
| # Citations | |
| if msg.get("citations"): | |
| story.append(Paragraph("📖 Sources:", citation_header_style)) | |
| for book_title, chapters in msg["citations"].items(): | |
| story.append(Paragraph(f"<b>{book_title}</b>", citation_style)) | |
| for ch in sorted(chapters): | |
| story.append(Paragraph(f" • {ch}", citation_style)) | |
| # Images | |
| if msg.get("images"): | |
| for img_data in msg["images"]: | |
| img_path = img_data.get("path","") | |
| if img_path and Path(img_path).exists(): | |
| try: | |
| rl_img = RLImage(img_path, width=12*cm, height=8*cm, kind='proportional') | |
| story.append(rl_img) | |
| story.append(Paragraph( | |
| f"Figure from: {img_data.get('chapter','')}", | |
| citation_style | |
| )) | |
| story.append(Spacer(1, 0.3*cm)) | |
| except Exception: | |
| pass | |
| story.append(Spacer(1, 0.2*cm)) | |
| story.append(Spacer(1, 0.5*cm)) | |
| story.append(HRFlowable(width="100%", thickness=1, color=colors.HexColor('#CCCCCC'))) | |
| story.append(Paragraph( | |
| "Generated by Remote Sensing RAG System", | |
| ParagraphStyle('Footer', parent=styles['Normal'], | |
| fontSize=8, textColor=colors.grey, | |
| alignment=TA_CENTER, spaceBefore=6) | |
| )) | |
| doc.build(story) | |
| buffer.seek(0) | |
| return buffer | |
| # ── LOAD RESOURCES ──────────────────────────────────────────── | |
| def load_resources(): | |
| data_path = Path(DB_LOCAL_PATH) | |
| # Download entire repo (DB + mineru output) on first startup | |
| if not data_path.exists() or not list(data_path.rglob("*.sqlite3")): | |
| st.info("First startup: downloading data (~may take 2-3 min)...") | |
| snapshot_download( | |
| repo_id=HF_DB_REPO, | |
| repo_type="dataset", | |
| local_dir=DB_LOCAL_PATH, | |
| ignore_patterns=["*.md", ".gitattributes"] | |
| ) | |
| chroma_client = chromadb.PersistentClient(path=DB_PATH) | |
| collection = chroma_client.get_or_create_collection(name=COLLECTION_NAME) | |
| embed_model = SentenceTransformer(EMBED_MODEL) | |
| hf_token = os.environ.get("HF_TOKEN", "") | |
| llm_client = InferenceClient( | |
| model=HF_LLM_MODEL, | |
| token=hf_token if hf_token else None | |
| ) | |
| return collection, embed_model, llm_client | |
| # ── PAGE CONFIG ─────────────────────────────────────────────── | |
| st.set_page_config( | |
| page_title="Remote Sensing RAG", | |
| page_icon="🛰️", | |
| layout="wide" | |
| ) | |
| st.title("🛰️ Remote Sensing RAG") | |
| st.caption("Retrieval-Augmented Generation over 6 Remote Sensing Textbooks") | |
| try: | |
| collection, embed_model, llm_client = load_resources() | |
| except Exception as e: | |
| st.error(f"Failed to load: {e}") | |
| st.stop() | |
| # ── SIDEBAR ─────────────────────────────────────────────────── | |
| with st.sidebar: | |
| st.header("⚙️ Settings") | |
| top_k = st.slider("Chunks to retrieve", 1, 10, TOP_K) | |
| threshold = st.slider("L2 Distance Threshold", 0.5, 2.0, L2_THRESHOLD, step=0.05) | |
| show_chunks = st.toggle("Show retrieved chunks", value=False) | |
| show_images = st.toggle("Show retrieved images", value=True) | |
| st.divider() | |
| # PDF Export | |
| st.header("💾 Export") | |
| if st.button("📄 Save chat as PDF", use_container_width=True): | |
| if not st.session_state.get("messages"): | |
| st.warning("No messages to export yet.") | |
| else: | |
| with st.spinner("Generating PDF..."): | |
| pdf_buffer = generate_chat_pdf(st.session_state.messages) | |
| st.download_button( | |
| label="⬇️ Download PDF", | |
| data=pdf_buffer, | |
| file_name=f"rs_rag_chat_{datetime.now().strftime('%Y%m%d_%H%M')}.pdf", | |
| mime="application/pdf", | |
| use_container_width=True | |
| ) | |
| st.divider() | |
| # DB Stats | |
| st.header("📚 Indexed Books") | |
| try: | |
| all_meta = collection.get(include=["metadatas"]) | |
| books = {} | |
| for meta in all_meta['metadatas']: | |
| ch_key = meta.get('chapter_name','') or Path(meta.get('source_file','')).stem | |
| book_title, _ = get_citation(meta.get('source_file',''), ch_key) | |
| books[book_title] = books.get(book_title, 0) + 1 | |
| for book, count in sorted(books.items()): | |
| st.markdown(f"**{book}**") | |
| st.caption(f"{count} chunks") | |
| except Exception as e: | |
| st.error(f"DB Error: {e}") | |
| st.divider() | |
| if st.button("🗑️ Clear chat", use_container_width=True): | |
| st.session_state.messages = [] | |
| st.rerun() | |
| st.caption(f"LLM: {HF_LLM_MODEL.split('/')[-1]}") | |
| # ── CHAT HISTORY ────────────────────────────────────────────── | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [] | |
| for message in st.session_state.messages: | |
| with st.chat_message(message["role"]): | |
| if message["role"] == "assistant": | |
| render_answer(message["content"]) | |
| else: | |
| st.markdown(message["content"]) | |
| if "citations" in message and message["citations"]: | |
| with st.expander("📖 Citations"): | |
| for book_title, chapters in message["citations"].items(): | |
| st.markdown(f"**{book_title}**") | |
| for ch in sorted(chapters): | |
| st.markdown(f" • {ch}") | |
| if "images" in message and show_images and message.get("images"): | |
| imgs = message["images"] | |
| if imgs: | |
| with st.expander(f"🖼️ Retrieved Images ({len(imgs)})"): | |
| cols = st.columns(min(len(imgs), 3)) | |
| for idx, img in enumerate(imgs): | |
| with cols[idx % 3]: | |
| try: | |
| st.image(img["path"], caption=img.get("chapter","")) | |
| except Exception: | |
| st.caption(f"Image unavailable") | |
| if "chunks" in message and show_chunks and message.get("chunks"): | |
| with st.expander(f"🔍 Retrieved Chunks ({len(message['chunks'])})"): | |
| for chunk in message["chunks"]: | |
| st.markdown(f"**{chunk['book']}** — {chunk['chapter']} | L2: {chunk['distance']}") | |
| st.text(chunk['text'][:300] + "...") | |
| st.divider() | |
| # ── QUERY ───────────────────────────────────────────────────── | |
| if query := st.chat_input("Ask a remote sensing question..."): | |
| st.session_state.messages.append({"role": "user", "content": query}) | |
| with st.chat_message("user"): | |
| st.markdown(query) | |
| with st.chat_message("assistant"): | |
| with st.spinner("Searching knowledge base..."): | |
| query_embedding = embed_model.encode(query).tolist() | |
| results = collection.query( | |
| query_embeddings=[query_embedding], | |
| n_results=top_k, | |
| include=["documents", "metadatas", "distances"] | |
| ) | |
| docs = results['documents'][0] | |
| metas = results['metadatas'][0] | |
| distances = results['distances'][0] | |
| filtered_docs = [] | |
| filtered_metas = [] | |
| retrieved_chunks = [] | |
| retrieved_images = [] | |
| for doc, meta, dist in zip(docs, metas, distances): | |
| if dist > threshold: | |
| continue | |
| ch_key = meta.get('chapter_name','') or Path(meta.get('source_file','')).stem | |
| book_title, chapter_title = get_citation(meta.get('source_file',''), ch_key) | |
| block_type = meta.get('block_type','text') | |
| filtered_docs.append(doc) | |
| filtered_metas.append(meta) | |
| retrieved_chunks.append({ | |
| "book": book_title, | |
| "chapter": chapter_title, | |
| "distance": round(dist, 3), | |
| "text": doc, | |
| }) | |
| # Image retrieval using the smart finder | |
| if block_type == 'image' and show_images: | |
| img_path_meta = meta.get('image_path','') | |
| found_path = find_image_in_mineru(img_path_meta, ch_key) | |
| if found_path: | |
| retrieved_images.append({ | |
| "path": found_path, | |
| "book": book_title, | |
| "chapter": chapter_title | |
| }) | |
| if not filtered_docs: | |
| answer = ( | |
| "Not enough context to answer this question from the indexed books. " | |
| "Try rephrasing, or check if this topic is covered in your textbooks." | |
| ) | |
| st.warning(answer) | |
| st.session_state.messages.append({"role": "assistant", "content": answer}) | |
| else: | |
| context = "\n\n".join(filtered_docs) | |
| # Last 1 exchange (2 messages) for context continuity | |
| history = "" | |
| recent = st.session_state.messages[-2:] if len(st.session_state.messages) >= 2 \ | |
| else st.session_state.messages | |
| for msg in recent: | |
| role = "User" if msg["role"] == "user" else "Assistant" | |
| history += f"{role}: {msg['content'][:500]}\n\n" | |
| prompt = f"""<s>[INST] You are a helpful remote sensing and geospatial analysis assistant. | |
| Use the retrieved context to answer the question clearly and in detail. | |
| If the answer is not in the context, say so briefly. | |
| Retrieved Context: | |
| {context} | |
| Conversation History: | |
| {history} | |
| Current Question: {query} [/INST]""" | |
| with st.spinner("Generating answer..."): | |
| try: | |
| answer = llm_client.text_generation( | |
| prompt, | |
| max_new_tokens=800, | |
| temperature=0.7, | |
| repetition_penalty=1.1, | |
| do_sample=True, | |
| ) | |
| render_answer(answer) | |
| # Citations | |
| citations = {} | |
| for meta in filtered_metas: | |
| ch_key = meta.get('chapter_name','') or Path(meta.get('source_file','')).stem | |
| book_title, chapter_title = get_citation(meta.get('source_file',''), ch_key) | |
| if book_title not in citations: | |
| citations[book_title] = set() | |
| if chapter_title: | |
| citations[book_title].add(chapter_title) | |
| with st.expander("📖 Citations"): | |
| for book_title, chapters in citations.items(): | |
| st.markdown(f"**{book_title}**") | |
| for ch in sorted(chapters): | |
| st.markdown(f" • {ch}") | |
| if retrieved_images and show_images: | |
| with st.expander(f"🖼️ Retrieved Images ({len(retrieved_images)})"): | |
| cols = st.columns(min(len(retrieved_images), 3)) | |
| for idx, img in enumerate(retrieved_images): | |
| with cols[idx % 3]: | |
| try: | |
| st.image(img["path"], caption=img["chapter"]) | |
| except Exception: | |
| st.caption("Image unavailable") | |
| if retrieved_chunks and show_chunks: | |
| with st.expander(f"🔍 Retrieved Chunks ({len(retrieved_chunks)})"): | |
| for chunk in retrieved_chunks: | |
| st.markdown(f"**{chunk['book']}** — {chunk['chapter']} | L2: {chunk['distance']}") | |
| st.text(chunk['text'][:300] + "...") | |
| st.divider() | |
| st.session_state.messages.append({ | |
| "role": "assistant", | |
| "content": answer, | |
| "citations": {k: list(v) for k, v in citations.items()}, | |
| "chunks": retrieved_chunks, | |
| "images": retrieved_images, | |
| }) | |
| except Exception as e: | |
| st.error(f"LLM Error: {e}") | |
| st.info("Free HF Inference API may be rate-limited. Wait a moment and retry.") |