HuggingFaceM4/obelics_visualization

1import json2 3import streamlit as st4from datasets import load_from_disk5 6 7class Visualization:8    def __init__(self, path_web_documents_dataset):9        self.path_web_documents_dataset = path_web_documents_dataset10 11    def visualization(self):12        self.set_title()13        self.load_dataset()14        self.choose_document()15        self.display_document()16 17    def set_title(self):18        st.title("Visualization of OBELICS web documents")19 20    def load_dataset(self):21        self.dataset = load_from_disk(self.path_web_documents_dataset)22 23    def choose_document(self):24        st.header("Choose a document")25        idx = st.number_input(26            f"Select a document among the first {self.dataset.num_rows} ones",27            min_value=0,28            max_value=self.dataset.num_rows - 1,29            value=0,30            step=1,31            help=f"Index between 0 and {self.dataset.num_rows-1}",32        )33        self.current_doc = self.dataset[idx]34 35    def display_document(self):36        st.header("Document")37        texts = self.current_doc["texts"]38        images = self.current_doc["images"]39        metadata = json.loads(self.current_doc["metadata"])40        for text, image, meta in zip(texts, images, metadata):41            if text:42                display_text = f"{text}\n".replace("\n", "<br>") # .replace(" ", "&nbsp;") Preserves white spaces, but creates text outside the width of the window43                st.markdown(f"<pre>{display_text}</pre>", unsafe_allow_html=True)44            elif image:45                st.markdown(f'<img src="{meta["src"]}" style="max-width: 1000px; height: auto;" />', unsafe_allow_html=True)46                st.text("\n")47                48 49if __name__ == "__main__":50    st.set_page_config(layout="wide")51    path_web_documents_dataset = "./web_docs_final_replaceimgbyurl"52    visualization = Visualization(path_web_documents_dataset=path_web_documents_dataset)53    visualization.visualization()54