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 | |
| # Constants | |
| CHUNK_SIZE = 32000 | |
| 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. | |
| """ | |
| # Initialize the Mistral chat model | |
| client = InferenceClient("mistralai/Mistral-Nemo-Instruct-2407") | |
| 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 extract_text_from_docx(docx_data, strip_content): | |
| """Extracts text from a DOCX file.""" | |
| text = u'' | |
| zipf = zipfile.ZipFile(io.BytesIO(docx_data)) | |
| filelist = zipf.namelist() | |
| for fname in filelist: | |
| if re.match('word/header[0-9]*.xml', fname): | |
| text += xml2text(zipf.read(fname)) | |
| elif re.match('word/footer[0-9]*.xml', fname): | |
| text += xml2text(zipf.read(fname)) | |
| text += xml2text(zipf.read('word/document.xml')) | |
| zipf.close() | |
| if strip_content: | |
| text = strip_text(text) | |
| return f"{text}\n\n**Document Length:** {len(text)} characters" | |
| def strip_text(text): | |
| """Strips unnecessary characters from text.""" | |
| content = text.replace('\n', ' ') | |
| content = content.replace('\r', ' ') | |
| content = content.replace('\t', ' ') | |
| content = content.replace(' ', '') | |
| return content.strip() | |
| def read_document(file, strip_content): | |
| """Reads the content of a document based on its file type.""" | |
| file_path = file.name | |
| file_extension = file_path.split('.')[-1].lower() | |
| with open(file_path, "rb") as f: | |
| file_content = f.read() | |
| if file_extension == 'pdf': | |
| 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 strip_content: | |
| content = strip_text(content) | |
| return content | |
| except Exception as e: | |
| return f"Error reading PDF: {e}" | |
| elif file_extension == 'xlsx': | |
| 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 strip_content: | |
| content = strip_text(content) | |
| return content | |
| except Exception as e: | |
| return f"Error reading XLSX: {e}" | |
| elif file_extension == 'pptx': | |
| try: | |
| presentation = Presentation(io.BytesIO(file_content)) | |
| content = '' | |
| for slide in presentation.slides: | |
| for shape in slide.shapes: | |
| if hasattr(shape, "text"): | |
| content += shape.text + ' ' | |
| if strip_content: | |
| content = strip_text(content) | |
| return content | |
| except Exception as e: | |
| return f"Error reading PPTX: {e}" | |
| elif file_extension == 'doc' or file_extension == 'docx': | |
| try: | |
| return extract_text_from_docx(file_content, strip_content) | |
| except Exception as e: | |
| return f"Error reading DOC/DOCX: {e}" | |
| else: | |
| try: | |
| content = file_content.decode('utf-8') | |
| if strip_content: | |
| content = strip_text(content) | |
| return content | |
| except Exception as e: | |
| return f"Error reading file: {e}" | |
| def split_content(content): | |
| """Splits content into chunks for processing.""" | |
| chunks = [] | |
| for i in range(0, len(content), CHUNK_SIZE): | |
| chunks.append(content[i:i + CHUNK_SIZE]) | |
| return chunks | |
| def chat_document(file, question, strip_content): | |
| """Handles chat with a document using Mistral.""" | |
| content = str(read_document(file, strip_content)) | |
| if len(content) > CHUNK_SIZE: | |
| content = content[:CHUNK_SIZE] | |
| message = f"""[INST] [SYSTEM] {SYSTEM_PROMPT} | |
| Document Content: {content} | |
| Question: {question} | |
| Answer:""" | |
| stream = client.text_generation(message, max_new_tokens=4096, 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_v2(file, question, strip_content): | |
| """Handles chat with a document using Mistral and chunk-based approach.""" | |
| content = str(read_document(file, strip_content)) | |
| chunks = split_content(content) | |
| all_answers = [] | |
| for chunk in chunks: | |
| message = f"""[INST] [SYSTEM] {SYSTEM_PROMPT} | |
| Document Content: {chunk[:CHUNK_SIZE]} | |
| Question: {question} | |
| Answer:""" | |
| stream = client.text_generation(message, max_new_tokens=4096, stream=True, details=True, return_full_text=False) | |
| output = "" | |
| for response in stream: | |
| if not response.token.text == "</s>": | |
| output += response.token.text | |
| all_answers.append(output) | |
| # 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) | |
| print(all_answers_str) | |
| summary_message = f"""[INST] [SYSTEM] {summary_prompt} | |
| {all_answers_str[:30000]} | |
| Summary:""" | |
| stream = client.text_generation(summary_message, max_new_tokens=4096, 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 | |
| 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="Strip Content", value=True)], | |
| outputs=gr.Textbox(label="Document Content"), | |
| 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="Strip Content", 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="Strip Content", 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() |