change to JSON output

app.py

@@ -7,11 +7,14 @@ try:
     api_key = st.secrets.hf_credentials.hf_api
 except:
     api_key = st.secrets.hf_api
-model_id = "meta-llama/Llama-2-13b-chat-hf"
+model_id = "meta-llama/Llama-2-70b-chat-hf"
 system_message = """
-    Your role is to take PDF documents and extract their raw text into a table format that can be uploaded into a database.
-    Return the table only. For example if you need to extract information about a report written on 2nd February 2011 with an author called Jane Mary then return this only: 
-    | report_written_date | author_name | \n | --- | --- | \n | 02/02/2011 | Jane Mary |
+    Your role is to take PDF documents and extract their raw text into a JSON format that can be uploaded into a database.
+    Return the JSON only. 
+    For example if you need to extract information about a report written on 2nd February 2011 with an author called Jane Mary then return this only: 
+    {'report_written_date': '02/02/2011', 'author_name': 'Jane Mary'}
+    Another example would be a clinical exam passed by a student on the 3rd of July 2022 would return this only:
+    {'result' : 'pass', 'date_of_exam' : '03/07/2022'}
     """
 
 
@@ -72,26 +75,35 @@ if st.session_state.key_inputs:
     keys_title = st.write("\nKeys/Columns for extraction:")
     keys_values = st.write(st.session_state.key_inputs)
 
-    if st.button("Extract data!"):
-        user_message = f"""
-            Use the text provided and denoted by 3 backticks ```{pdf_text}```. 
-            Extract the following columns and return a table that could be uploaded to an SQL database. 
-            {'; '.join([key + ': ' + st.session_state.key_inputs[key] for key in st.session_state.key_inputs])}
-        """
-        the_prompt = prompt_generator(
-            system_message=system_message, user_message=user_message
-        )
-        response = query(
-            {
-                "inputs": the_prompt,
-                "parameters": {"max_new_tokens": 500, "temperature": 0.1},
-            },
-            model_id,
-        )
-        match = re.search(
-            pattern, response[0]["generated_text"], re.MULTILINE | re.DOTALL
-        )
-        if match:
-            response = match.group(1).strip()
-
-        st.markdown(f"Data Extracted!\n{response}")
+    with st.spinner("Extracting requested data"):
+        if st.button("Extract data!"):
+            user_message = f"""
+                Use the text provided and denoted by 3 backticks ```{pdf_text}```. 
+                Extract the following columns and return a table that could be uploaded to an SQL database.
+                {'; '.join([key + ': ' + st.session_state.key_inputs[key] for key in st.session_state.key_inputs])}
+            """
+            the_prompt = prompt_generator(
+                system_message=system_message, user_message=user_message
+            )
+            response = query(
+                {
+                    "inputs": the_prompt,
+                    "parameters": {"max_new_tokens": 500, "temperature": 0.1},
+                },
+                model_id,
+            )
+            try:
+                match = re.search(
+                    pattern, response[0]["generated_text"], re.MULTILINE | re.DOTALL
+                )
+                if match:
+                    response = match.group(1).strip()
+
+                response = eval(response)
+
+                st.success("Data Extracted Successfully!")
+                st.write(response)
+            except:
+                st.error("Unable to connect to model. Please try again later.")
+
+            # st.success(f"Data Extracted!")