Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import os | |
| import pandas as pd | |
| import fitz # PyMuPDF | |
| from google import genai | |
| from google.genai import types | |
| from docx import Document | |
| from pptx import Presentation | |
| from reportlab.lib.pagesizes import A4 | |
| from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer | |
| from reportlab.lib.styles import getSampleStyleSheet | |
| from bs4 import BeautifulSoup | |
| import re | |
| import base64 | |
| # Set Gemini API Key | |
| client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY")) | |
| # --------- File Extractors --------- | |
| def extract_text_from_docx(docx_file): | |
| document = Document(docx_file) | |
| return "\n".join([para.text for para in document.paragraphs]) | |
| def extract_text_from_csv(csv_file): | |
| df = pd.read_csv(csv_file) | |
| return df.to_string(index=False) | |
| def extract_text_from_xlsx(xlsx_file): | |
| df = pd.read_excel(xlsx_file) | |
| return df.to_string(index=False) | |
| def extract_text_from_pptx(pptx_file): | |
| prs = Presentation(pptx_file) | |
| text_runs = [] | |
| for slide in prs.slides: | |
| for shape in slide.shapes: | |
| if hasattr(shape, "text"): | |
| text_runs.append(shape.text) | |
| return "\n".join(text_runs) | |
| def extract_text_from_pdf(pdf_file): | |
| doc = fitz.open(stream=pdf_file.read(), filetype="pdf") | |
| text = "" | |
| for page in doc: | |
| text += page.get_text() | |
| return text | |
| def extract_text_from_html(html_file): | |
| soup = BeautifulSoup(html_file.read(), "html.parser") | |
| return soup.get_text() | |
| def extract_text_from_tex(tex_file): | |
| content = tex_file.read().decode("utf-8") | |
| content = re.sub(r'\\[a-zA-Z]+\{[^}]*\}', '', content) | |
| content = re.sub(r'\\[a-zA-Z]+', '', content) | |
| return content | |
| def process_image(image_file): | |
| image_bytes = image_file.read() | |
| encoded_image = base64.b64encode(image_bytes).decode("utf-8") | |
| return { | |
| "inline_data": { | |
| "mime_type": image_file.type, | |
| "data": encoded_image | |
| } | |
| } | |
| def process_video(video_file): | |
| video_bytes = video_file.read() | |
| filetype = video_file.name.split(".")[-1].lower() | |
| mime_map = { | |
| "avi": "video/x-msvideo", | |
| "mkv": "video/x-matroska", | |
| "flv": "video/x-flv", | |
| "wmv": "video/x-ms-wmv", | |
| "mpeg": "video/mpeg" | |
| } | |
| mime_type = mime_map.get(filetype, video_file.type) | |
| encoded_video = base64.b64encode(video_bytes).decode("utf-8") | |
| return { | |
| "inline_data": { | |
| "mime_type": mime_type, | |
| "data": encoded_video | |
| } | |
| } | |
| def export_conversation_to_pdf(conversation_history): | |
| pdf_path = "Conversation.pdf" | |
| doc = SimpleDocTemplate(pdf_path, pagesize=A4) | |
| styles = getSampleStyleSheet() | |
| elements = [] | |
| for i, (user_q, gemini_a) in enumerate(conversation_history): | |
| elements.append(Paragraph(f"<b>Q{i+1}:</b> {user_q}", styles["Normal"])) | |
| elements.append(Spacer(1, 8)) | |
| elements.append(Paragraph(f"<b>A{i+1}:</b> {gemini_a.replace(chr(10), '<br/>')}", styles["Normal"])) | |
| elements.append(Spacer(1, 16)) | |
| doc.build(elements) | |
| return pdf_path | |
| # --------- Main App --------- | |
| def main(): | |
| st.set_page_config(page_title="Gemini 2.5 Flash Chat Q&A", layout="wide") | |
| st.title("\U0001F4C4 Product Prototype with Document Support and Internet Search") | |
| if "conversation" not in st.session_state: | |
| st.session_state.conversation = [] | |
| if "documents_text" not in st.session_state: | |
| st.session_state.documents_text = [] | |
| if "chat_active" not in st.session_state: | |
| st.session_state.chat_active = True | |
| if "chat_history" not in st.session_state: | |
| st.session_state.chat_history = [] | |
| uploaded_files = st.file_uploader( | |
| "Upload files (.txt, .docx, .csv, .xlsx, .pptx, .pdf, .html, .htm, .tex, .jpg, .jpeg, .png, video formats):", | |
| type=[ | |
| "txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex", | |
| "jpg", "jpeg", "png", | |
| "mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg" | |
| ], | |
| accept_multiple_files=True | |
| ) | |
| if uploaded_files: | |
| all_text_content = "" | |
| for uploaded_file in uploaded_files: | |
| filetype = uploaded_file.name.split(".")[-1].lower() | |
| try: | |
| if filetype == "txt": | |
| all_text_content += uploaded_file.read().decode("utf-8") + "\n" | |
| elif filetype == "docx": | |
| all_text_content += extract_text_from_docx(uploaded_file) + "\n" | |
| elif filetype == "csv": | |
| all_text_content += extract_text_from_csv(uploaded_file) + "\n" | |
| elif filetype == "xlsx": | |
| all_text_content += extract_text_from_xlsx(uploaded_file) + "\n" | |
| elif filetype == "pptx": | |
| all_text_content += extract_text_from_pptx(uploaded_file) + "\n" | |
| elif filetype == "pdf": | |
| all_text_content += extract_text_from_pdf(uploaded_file) + "\n" | |
| elif filetype in ["html", "htm"]: | |
| all_text_content += extract_text_from_html(uploaded_file) + "\n" | |
| elif filetype == "tex": | |
| all_text_content += extract_text_from_tex(uploaded_file) + "\n" | |
| elif filetype in ["jpg", "jpeg", "png"]: | |
| st.session_state.image_data = process_image(uploaded_file) | |
| all_text_content += "Image uploaded and processed.\n" | |
| st.image(uploaded_file, caption=uploaded_file.name, use_container_width=True) | |
| elif filetype in ["mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"]: | |
| st.session_state.video_data = process_video(uploaded_file) | |
| all_text_content += f"Video file '{uploaded_file.name}' uploaded and processed.\n" | |
| st.video(uploaded_file) | |
| else: | |
| st.error(f"Unsupported file format: {uploaded_file.name}") | |
| except Exception as e: | |
| st.error(f"Failed to extract/process {uploaded_file.name}: {e}") | |
| st.session_state.documents_text = all_text_content | |
| if st.session_state.documents_text: | |
| st.markdown("### \U0001F4AC Conversation") | |
| for user_q, gemini_a in st.session_state.conversation: | |
| st.markdown(f"**You:** {user_q}") | |
| st.markdown(f"**Gemini:**\n\n{gemini_a}") | |
| if st.session_state.chat_active: | |
| with st.form(key="chat_form", clear_on_submit=True): | |
| user_input = st.text_input("Ask a question about the documents (type 'exit' to stop):") | |
| submit = st.form_submit_button("Send") | |
| if submit and user_input: | |
| if user_input.strip().lower() == "exit": | |
| st.session_state.chat_active = False | |
| st.success("Chat ended. Reload to start again.") | |
| else: | |
| with st.spinner("Let me think..."): | |
| content_blocks = [] | |
| # System prompt inserted first here: | |
| content_blocks.append({ | |
| "text": ( | |
| "Your role is to analyze documents, images, and videos to help identify and classify driver states, " | |
| "including: drowsy, alert, yawning, microsleep, eyes/lips open or closed, and other distractions. " | |
| "The user may ask for timestamps of events in video, technical explanations from code/docs, " | |
| "or structured summaries. Prioritize your answers for this use case and act as a helpful assistant " | |
| "within this domain. " | |
| "Additionally, your classification may rely on the State Farm Distracted Driver Detection dataset, " | |
| "which defines 10 distraction categories such as texting, phone use, reaching behind, adjusting the radio, " | |
| "and safe driving. Use this labeling scheme when interpreting relevant images." | |
| ) | |
| }) | |
| if st.session_state.conversation: | |
| history_text = "\n\n".join( | |
| f"Q: {entry[0]}\nA: {entry[1]}" | |
| for entry in st.session_state.conversation | |
| ) | |
| content_blocks.append({"text": f"Previous conversation:\n{history_text}"}) | |
| if st.session_state.documents_text.strip(): | |
| content_blocks.append({"text": f"Context:\n{st.session_state.documents_text}"}) | |
| if "image_data" in st.session_state: | |
| content_blocks.append(st.session_state.image_data) | |
| if "video_data" in st.session_state: | |
| content_blocks.append(st.session_state.video_data) | |
| content_blocks.append({"text": f"Question: {user_input}"}) | |
| response = client.models.generate_content( | |
| model="gemini-2.5-flash", | |
| contents=content_blocks, | |
| config={"tools": [{"google_search": {}}]} | |
| ) | |
| st.session_state.conversation.append((user_input, response.text)) | |
| st.success("\U0001F4A1 Answer:") | |
| st.write(response.text) | |
| st.session_state.chat_history.append({ | |
| "question": user_input, | |
| "answer": response.text | |
| }) | |
| if st.session_state.conversation: | |
| pdf_path = export_conversation_to_pdf(st.session_state.conversation) | |
| with open(pdf_path, "rb") as f: | |
| st.download_button( | |
| label="\U0001F4E5 Export Conversation as PDF", | |
| data=f, | |
| file_name="Conversation.pdf", | |
| mime="application/pdf" | |
| ) | |
| if __name__ == "__main__": | |
| main() |