Spaces:
Running
Running
| import PyPDF2 | |
| from openpyxl import load_workbook | |
| from pptx import Presentation | |
| import gradio as gr | |
| import io | |
| from huggingface_hub import InferenceClient | |
| import re | |
| import zipfile | |
| import xml.etree.ElementTree as ET | |
| import filetype | |
| # Constants | |
| CHUNK_SIZE = 32000 | |
| MAX_NEW_TOKENS = 4096 | |
| # Initialize the Mistral chat model | |
| client = InferenceClient("mistralai/Mistral-Nemo-Instruct-2407") | |
| # --- Utility Functions --- | |
| def xml2text(xml): | |
| """Extracts text from XML data.""" | |
| text = u'' | |
| root = ET.fromstring(xml) | |
| for child in root.iter(): | |
| text += child.text + " " if child.text is not None else '' | |
| return text | |
| def clean_text(content): | |
| """Cleans text content based on the 'clean' parameter.""" | |
| content = content.replace('\n', ' ') | |
| content = content.replace('\r', ' ') | |
| content = content.replace('\t', ' ') | |
| content = re.sub(r'\s+', ' ', content) | |
| return content | |
| def split_content(content, chunk_size=CHUNK_SIZE): | |
| """Splits content into chunks of a specified size.""" | |
| chunks = [] | |
| for i in range(0, len(content), chunk_size): | |
| chunks.append(content[i:i + chunk_size]) | |
| return chunks | |
| # --- Document Reading Functions --- | |
| def extract_text_from_docx(docx_data, clean=True): | |
| """Extracts text from DOCX files.""" | |
| text = u'' | |
| zipf = zipfile.ZipFile(io.BytesIO(docx_data)) | |
| filelist = zipf.namelist() | |
| header_xmls = 'word/header[0-9]*.xml' | |
| for fname in filelist: | |
| if re.match(header_xmls, fname): | |
| text += xml2text(zipf.read(fname)) | |
| doc_xml = 'word/document.xml' | |
| text += xml2text(zipf.read(doc_xml)) | |
| footer_xmls = 'word/footer[0-9]*.xml' | |
| for fname in filelist: | |
| if re.match(footer_xmls, fname): | |
| text += xml2text(zipf.read(fname)) | |
| zipf.close() | |
| if clean: | |
| text = clean_text(text) | |
| return text, len(text) | |
| def extract_text_from_pptx(pptx_data, clean=True): | |
| """Extracts text from PPT files.""" | |
| text = u'' | |
| zipf = zipfile.ZipFile(io.BytesIO(pptx_data)) | |
| filelist = zipf.namelist() | |
| # Extract text from slide notes | |
| notes_xmls = 'ppt/notesSlides/notesSlide[0-9]*.xml' | |
| for fname in filelist: | |
| if re.match(notes_xmls, fname): | |
| text += xml2text(zipf.read(fname)) | |
| # Extract text from slide content (shapes and text boxes) | |
| slide_xmls = 'ppt/slides/slide[0-9]*.xml' | |
| for fname in filelist: | |
| if re.match(slide_xmls, fname): | |
| text += xml2text(zipf.read(fname)) | |
| zipf.close() | |
| if clean: | |
| text = clean_text(text) | |
| return text, len(text) | |
| def read_document(file, clean=True): | |
| """Reads content from various document formats.""" | |
| file_path = file.name | |
| # No file extension used | |
| with open(file_path, "rb") as f: | |
| file_content = f.read() | |
| kind = filetype.guess(file_content) | |
| if kind is None: | |
| return "Cannot guess file type", 0 # Handle unknown file types | |
| mime = kind.mime | |
| if mime == "application/pdf": | |
| # PDF Handling (unchanged) | |
| try: | |
| pdf_reader = PyPDF2.PdfReader(io.BytesIO(file_content)) | |
| content = '' | |
| for page in range(len(pdf_reader.pages)): | |
| content += pdf_reader.pages[page].extract_text() | |
| if clean: | |
| content = clean_text(content) | |
| return content, len(content) | |
| except Exception as e: | |
| return f"Error reading PDF: {e}", 0 | |
| elif mime == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": | |
| # XLSX Handling (unchanged) | |
| try: | |
| wb = load_workbook(io.BytesIO(file_content)) | |
| content = '' | |
| for sheet in wb.worksheets: | |
| for row in sheet.rows: | |
| for cell in row: | |
| if cell.value is not None: | |
| content += str(cell.value) + ' ' | |
| if clean: | |
| content = clean_text(content) | |
| return content, len(content) | |
| except Exception as e: | |
| return f"Error reading XLSX: {e}", 0 | |
| elif mime == "text/plain": | |
| try: | |
| content = file_content.decode('utf-8') | |
| if clean: | |
| content = clean_text(content) | |
| return content, len(content) | |
| except Exception as e: | |
| return f"Error reading TXT file: {e}", 0 | |
| elif mime == "text/csv": | |
| try: | |
| content = file_content.decode('utf-8') | |
| if clean: | |
| content = clean_text(content) | |
| return content, len(content) | |
| except Exception as e: | |
| return f"Error reading CSV file: {e}", 0 | |
| elif mime == "application/vnd.openxmlformats-officedocument.wordprocessingml.document": | |
| try: | |
| return extract_text_from_docx(file_content, clean) | |
| except Exception as e: | |
| return f"Error reading DOCX: {e}", 0 | |
| elif mime == "application/vnd.openxmlformats-officedocument.presentationml.presentation": | |
| try: | |
| return extract_text_from_pptx(file_content, clean) | |
| except Exception as e: | |
| return f"Error reading PPTX: {e}", 0 | |
| else: | |
| try: | |
| content = file_content.decode('utf-8') | |
| if clean: | |
| content = clean_text(content) | |
| return content, len(content) | |
| except Exception as e: | |
| return f"Error reading file: {e}", 0 | |
| # --- Chat Functions --- | |
| def generate_mistral_response(message): | |
| """Generates a response from the Mistral API.""" | |
| stream = client.text_generation( | |
| message, | |
| max_new_tokens=MAX_NEW_TOKENS, | |
| stream=True, | |
| details=True, | |
| return_full_text=False | |
| ) | |
| output = "" | |
| for response in stream: | |
| if not response.token.text == "</s>": | |
| output += response.token.text | |
| yield output | |
| def chat_document(file, question, clean=True): | |
| """Chats with a document using a single Mistral API call.""" | |
| content, length = read_document(file, clean) | |
| if length > CHUNK_SIZE: | |
| content = content[:CHUNK_SIZE] # Limit to max chunk size | |
| system_prompt = """ | |
| You are a helpful and informative assistant that can answer questions based on the content of documents. | |
| You will receive the content of a document and a question about it. | |
| Your task is to provide a concise and accurate answer to the question based solely on the provided document content. | |
| If the document does not contain enough information to answer the question, simply state that you cannot answer the question based on the provided information. | |
| """ | |
| message = f"""[INST] [SYSTEM] {system_prompt} | |
| Document Content: {content} | |
| Question: {question} | |
| Answer:""" | |
| yield from generate_mistral_response(message) | |
| def chat_document_v2(file, question, clean=True): | |
| """Chats with a document using chunk-based Mistral API calls and summarizes the answers.""" | |
| content, length = read_document(file, clean) | |
| chunks = split_content(content) | |
| system_prompt = """ | |
| You are a helpful and informative assistant that can answer questions based on the content of documents. | |
| You will receive the content of a document and a question about it. | |
| Your task is to provide a concise and accurate answer to the question based solely on the provided document content. | |
| If the document does not contain enough information to answer the question, simply state that you cannot answer the question based on the provided information. | |
| """ | |
| all_answers = [] | |
| for chunk in chunks: | |
| message = f"""[INST] [SYSTEM] {system_prompt} | |
| Document Content: {chunk[:CHUNK_SIZE]} | |
| Question: {question} | |
| Answer:""" | |
| response = "" | |
| for stream_response in generate_mistral_response(message): | |
| response = stream_response # Update with latest response | |
| all_answers.append(response) | |
| # Summarize all answers using Mistral | |
| summary_prompt = """ | |
| You are a helpful and informative assistant that can summarize multiple answers related to the same question. | |
| You will receive a list of answers to a question, and your task is to generate a concise and comprehensive summary that incorporates the key information from all the answers. | |
| Avoid repeating information unnecessarily and focus on providing the most relevant and accurate summary based on the provided answers. | |
| Answers: | |
| """ | |
| all_answers_str = "\n".join(all_answers) | |
| summary_message = f"""[INST] [SYSTEM] {summary_prompt} | |
| {all_answers_str[:30000]} | |
| Summary:""" | |
| yield from generate_mistral_response(summary_message) | |
| # --- Gradio Interface --- | |
| with gr.Blocks() as demo: | |
| with gr.Tabs(): | |
| with gr.TabItem("Document Reader"): | |
| iface1 = gr.Interface( | |
| fn=read_document, | |
| inputs=[ | |
| gr.File(label="Upload a Document"), | |
| gr.Checkbox(label="Clean Text", value=True), | |
| ], | |
| outputs=[ | |
| gr.Textbox(label="Document Content"), | |
| gr.Number(label="Document Length (characters)"), | |
| ], | |
| title="Document Reader", | |
| description="Upload a document (PDF, XLSX, PPTX, TXT, CSV, DOC, DOCX and Code or text file) to read its content." | |
| ) | |
| with gr.TabItem("Document Chat"): | |
| iface2 = gr.Interface( | |
| fn=chat_document, | |
| inputs=[ | |
| gr.File(label="Upload a Document"), | |
| gr.Textbox(label="Question"), | |
| gr.Checkbox(label="Clean and Compress Text", value=True), | |
| ], | |
| outputs=gr.Markdown(label="Answer"), | |
| title="Document Chat", | |
| description="Upload a document and ask questions about its content." | |
| ) | |
| with gr.TabItem("Document Chat V2"): | |
| iface3 = gr.Interface( | |
| fn=chat_document_v2, | |
| inputs=[ | |
| gr.File(label="Upload a Document"), | |
| gr.Textbox(label="Question"), | |
| gr.Checkbox(label="Clean Text", value=True), | |
| ], | |
| outputs=gr.Markdown(label="Answer"), | |
| title="Document Chat V2", | |
| description="Upload a document and ask questions about its content (using chunk-based approach)." | |
| ) | |
| demo.launch() |