DeepakJaiz/Extracting_term_defination

1import os2import streamlit as st3 4from PIL import Image5from llama_index import (6    Document,7    GPTSimpleVectorIndex,8    GPTListIndex,9    LLMPredictor,10    ServiceContext,11    SimpleDirectoryReader,12    PromptHelper,13  14 15)16from llama_index.readers.file.base import (DEFAULT_FILE_EXTRACTOR, ImageParser)17 18from constants import DEFAULT_TERM_STR, DEFAULT_TERMS, REFINE_TEMPLATE, TEXT_QA_TEMPLATE19from utils import get_llm20 21 22if "all_terms" not in st.session_state:23    st.session_state["all_terms"] = DEFAULT_TERMS24 25 26@st.cache_resource27def get_file_extractor():28    image_parser = ImageParser(keep_image=True, parse_text=True)29    file_extractor = DEFAULT_FILE_EXTRACTOR30    file_extractor.update(31        {32            ".jpg": image_parser,33            ".png": image_parser,34            ".jpeg": image_parser,35        }36    )37 38    return file_extractor39 40 41file_extractor = get_file_extractor()42 43 44def extract_terms(documents, term_extract_str, llm_name, model_temperature, api_key):45    llm = get_llm(llm_name, model_temperature, api_key, max_tokens=1024)46 47    service_context = ServiceContext.from_defaults(48        llm_predictor=LLMPredictor(llm=llm),49        prompt_helper=PromptHelper(50            max_input_size=4096, max_chunk_overlap=20, num_output=102451        ),52        chunk_size_limit=1024,53    )54 55    temp_index = GPTListIndex.from_documents(documents, service_context=service_context)56    terms_definitions = str(57        temp_index.query(term_extract_str, response_mode="tree_summarize")58    )59    terms_definitions = [60        x61        for x in terms_definitions.split("\n")62        if x and "Term:" in x and "Definition:" in x63    ]64    # parse the text into a dict65    terms_to_definition = {66        x.split("Definition:")[0]67        .split("Term:")[-1]68        .strip(): x.split("Definition:")[-1]69        .strip()70        for x in terms_definitions71    }72    return terms_to_definition73 74 75def insert_terms(terms_to_definition):76    for term, definition in terms_to_definition.items():77        doc = Document(f"Term: {term}\nDefinition: {definition}")78        st.session_state["llama_index"].insert(doc)79 80@st.cache_resource81def initialize_index(llm_name, model_temperature, api_key):82    """Create the GPTSQLStructStoreIndex object."""83    llm = get_llm(llm_name, model_temperature, api_key)84 85    service_context = ServiceContext.from_defaults(llm_predictor=LLMPredictor(llm=llm))86 87    index = GPTSimpleVectorIndex.load_from_disk(88        "./index.json", service_context=service_context89    )90 91    return index92 93 94st.title("🦙 Llama Index Term Extractor 🦙")95st.markdown(96    (97        "This demo allows you to upload your own documents (either a screenshot/image or the actual text) and extract terms and definitions, building a knowledge base!\n\n"98        "Powered by [Llama Index](https://gpt-index.readthedocs.io/en/latest/index.html) and OpenAI, you can augment the existing knowledge of an "99        "LLM using your own notes, documents, and images. Then, when you ask about a term or definition, it will use your data first! "100        "The app is currently pre-loaded with terms from the NYC Wikipedia page."101    )102)103 104setup_tab, terms_tab, upload_tab, query_tab = st.tabs(105    ["Setup", "All Terms", "Upload/Extract Terms", "Query Terms"]106)107 108with setup_tab:109    st.subheader("LLM Setup")110    api_key = st.text_input("Enter your OpenAI API key here", type="password")111    llm_name = st.selectbox(112        "Which LLM?", ["text-davinci-003", "gpt-3.5-turbo", "gpt-4"]113    )114    model_temperature = st.slider(115        "LLM Temperature", min_value=0.0, max_value=1.0, step=0.1116    )117    term_extract_str = st.text_area(118        "The query to extract terms and definitions with.", value=DEFAULT_TERM_STR119    )120 121 122with terms_tab:123    st.subheader("Current Extracted Terms and Definitions")124    st.json(st.session_state["all_terms"])125 126 127with upload_tab:128    st.subheader("Extract and Query Definitions")129    if st.button("Initialize Index and Reset Terms", key="init_index_1"):130        st.session_state["llama_index"] = initialize_index(131            llm_name, model_temperature, api_key132        )133        st.session_state["all_terms"] = DEFAULT_TERMS134 135    if "llama_index" in st.session_state:136        st.markdown(137            "Either upload an image/screenshot of a document, or enter the text manually."138        )139        uploaded_file = st.file_uploader(140            "Upload an image/screenshot of a document:", type=["png", "jpg", "jpeg"]141        )142        document_text = st.text_area("Or enter raw text")143        if st.button("Extract Terms and Definitions") and (144            uploaded_file or document_text145        ):146            st.session_state["terms"] = {}147            terms_docs = {}148            with st.spinner("Extracting (images may be slow)..."):149                if document_text:150                    terms_docs.update(151                        extract_terms(152                            [Document(document_text)],153                            term_extract_str,154                            llm_name,155                            model_temperature,156                            api_key,157                        )158                    )159                if uploaded_file:160                    Image.open(uploaded_file).convert("RGB").save("temp.png")161                    img_reader = SimpleDirectoryReader(162                        input_files=["temp.png"], file_extractor=file_extractor163                    )164                    img_docs = img_reader.load_data()165                    os.remove("temp.png")166                    terms_docs.update(167                        extract_terms(168                            img_docs,169                            term_extract_str,170                            llm_name,171                            model_temperature,172                            api_key,173                        )174                    )175            st.session_state["terms"].update(terms_docs)176 177    if "terms" in st.session_state and st.session_state["terms"]:178        st.markdown("Extracted terms")179        st.json(st.session_state["terms"])180 181        if st.button("Insert terms?"):182            with st.spinner("Inserting terms"):183                insert_terms(st.session_state["terms"])184            st.session_state["all_terms"].update(st.session_state["terms"])185            st.session_state["terms"] = {}186            st.experimental_rerun()187 188with query_tab:189    st.subheader("Query for Terms/Definitions!")190    st.markdown(191        (192            "The LLM will attempt to answer your query, and augment it's answers using the terms/definitions you've inserted. "193            "If a term is not in the index, it will answer using it's internal knowledge."194        )195    )196    if st.button("Initialize Index and Reset Terms", key="init_index_2"):197        st.session_state["llama_index"] = initialize_index(198            llm_name, model_temperature, api_key199        )200        st.session_state["all_terms"] = DEFAULT_TERMS201 202    if "llama_index" in st.session_state:203        query_text = st.text_input("Ask about a term or definition:")204        if query_text:205            with st.spinner("Generating answer..."):206                response = st.session_state["llama_index"].query(207                    query_text, similarity_top_k=5, response_mode="compact",208                    text_qa_template=TEXT_QA_TEMPLATE, refine_template=REFINE_TEMPLATE209                )210            st.markdown(str(response))