Update app.py

app.py

@@ -9,9 +9,10 @@ from langchain.prompts.prompt import PromptTemplate
 from langchain.embeddings import HuggingFaceEmbeddings
 import json
 import hashlib
+import pdf2bib
 # Initialize the storage directories if they don't exist
 model_name = "intfloat/e5-large-v2"
-model_kwargs = {'device': 'cpu'}
+model_kwargs = {'device': 'cuda'}
 encode_kwargs = {'normalize_embeddings': False}
 embeddings = HuggingFaceEmbeddings(
     model_name=model_name,
@@ -73,13 +74,14 @@ def save_metadata(metadata, metadata_dir):
         json.dump(metadata, f, indent=4)
 
 #Add a new entry to the metadata
-def add_metadata(file_path, title, author, year, publisher, DOI,page,volume,issue):
+def add_metadata(file_path, title, author, year,journal, publisher, DOI,page,volume,issue):
      metadata = load_metadata(st.session_state['metadata_dir'])
      metadata_dir = st.session_state['metadata_dir']
      metadata[file_path] = {
          'title': title,
          'author': author,
          'year': year,
+         'journal': journal,
          'publisher': publisher,
          'DOI': DOI,
          'page':page,
@@ -112,6 +114,7 @@ def edit_metadata_of_selected_pdf(metadata):
                 edited_volume = st.text_input('volume', value=data['volume'])
 
             with col2:
+                edited_journal = st.text_input('journal', value=data['journal'])
                 edited_publisher = st.text_input('Publisher', value=data['publisher'])
                 edited_DOI = st.text_input('DOI', value=data['DOI'])
                 edited_page = st.text_input('Page', value=data['page'])
@@ -126,6 +129,7 @@ def edit_metadata_of_selected_pdf(metadata):
                     'title': edited_title,
                     'author': edited_author,
                     'year': edited_year,
+                    'journal': edited_journal,
                     'publisher': edited_publisher,
                     'DOI': edited_DOI,
                     'page': edited_page,
@@ -140,13 +144,15 @@ def edit_metadata_of_selected_pdf(metadata):
 # Function to delete selected PDF and its metadata
 def delete_pdf_and_metadata(metadata):
     # Use a select box for users to choose a PDF to delete
+    metadata_dir = st.session_state['metadata_dir']
+    st.write(list(metadata.keys()))
     file_names = list(metadata.keys())
     if not file_names:
         st.write("No PDFs available to delete.")
         return
 
     selected_file_to_delete = st.selectbox('Select a PDF file to delete', file_names)
-
+    st.write(selected_file_to_delete)
     # Button to delete the PDF and its metadata
     if st.button(f"Delete '{selected_file_to_delete}' and its metadata"):
         # Delete the PDF file
@@ -154,10 +160,10 @@ def delete_pdf_and_metadata(metadata):
 
         # Delete the metadata entry
         del metadata[selected_file_to_delete]
-        save_metadata(metadata)
+        save_metadata(metadata, metadata_dir)
 
         # Update the display
-        st.experimental_rerun()
+        #st.experimental_rerun()
 
 def generate_response(input_text,openai_api_key):
     db = FAISS.load_local(st.session_state['embed_dir'], embeddings)
@@ -172,14 +178,14 @@ def generate_response(input_text,openai_api_key):
                         {context}
                         Question: {question}
 
-                        use metadata in the "metadata" html block to create APA style references and list them in the reference section:
+                        use metadata in the "metadata" html block to create APA style references and list them in a reference section:
                         <metadata>
                         {source1},{source2},{source3},{source4},{source5}
                         <metadata/>
 
                         make sure to add references
 
-                        format the answer in mardown format
+                        format the answer in markdown format
 
                         """
     QA_PROMPT = PromptTemplate(input_variables=["question", "context"],
@@ -233,6 +239,20 @@ def init_user_assets(username):
         os.makedirs(embed_dir)
     return assets_dir, pdf_dir, metadata_dir, embed_dir
 
+
+
+def all_author_names(data):
+    author_full_names = []
+    authors_string =[]
+# Iterate over each author in the author list
+    for author in data["author"]:
+        # Combine given and family names
+        full_name = f"{author['family']},{author['given']}"
+        author_full_names.append(full_name)
+
+# Join the list of author names into a single string separated by commas
+    authors_string = "; ".join(author_full_names)
+    return authors_string
 # Main App
 def main():
     st.title('LLM powered paper management system')
@@ -285,24 +305,33 @@ def main():
                         if uploaded_file is not None:
                             # Save file
                             file_path = os.path.join(st.session_state['pdf_dir'], uploaded_file.name)
+
                             with open(file_path, "wb") as f:
                                 f.write(uploaded_file.getbuffer())
-
+                            pdfextractdata = pdf2bib.pdf2bib(file_path)
+                            #st.write(pdfextractdata['metadata'])
                             # Collect metadata
                             with st.form(key=uploaded_file.name):
                                 st.write("Metadata for:", uploaded_file.name)
-                                title = st.text_input('Title', key=f'title_{uploaded_file.name}')
-                                author = st.text_input('Author', key=f'author_{uploaded_file.name}')
-                                year = st.text_input('Published Year', key=f'year_{uploaded_file.name}')
-                                publisher = st.text_input('Publisher', key=f'publisher_{uploaded_file.name}')
-                                DOI = st.text_input('DOI', key=f'DOI_{uploaded_file.name}')
-                                page = st.text_input('page', key=f'page_{uploaded_file.name}')
-                                volume = st.text_input('volume', key=f'volume_{uploaded_file.name}')
-                                issue = st.text_input('issue', key=f'issue_{uploaded_file.name}')
+                                col1, col2, col3 = st.columns(3)
+
+                                with col1:
+                                    title = st.text_input('Title', key=f'title_{uploaded_file.name}', value=pdfextractdata['metadata'].get('title', ''))
+                                    year = st.text_input('Published Year', key=f'year_{uploaded_file.name}', value=pdfextractdata['metadata'].get('year', ''))
+                                    volume = st.text_input('Volume', key=f'volume_{uploaded_file.name}', value=pdfextractdata['metadata'].get('volume', ''))
+
+                                with col2:
+                                    author = st.text_input('Author', key=f'author_{uploaded_file.name}', value=all_author_names(pdfextractdata['metadata']))
+                                    journal = st.text_input('Journal', key=f'journal_{uploaded_file.name}', value=pdfextractdata['metadata'].get('journal', ''))
+                                    issue = st.text_input('Issue', key=f'issue_{uploaded_file.name}', value=pdfextractdata['metadata'].get('issue', ''))
+
+                                with col3:
+                                    publisher = st.text_input('Publisher', key=f'publisher_{uploaded_file.name}', value=pdfextractdata['metadata'].get('publisher', ''))
+                                    DOI = st.text_input('DOI', key=f'DOI_{uploaded_file.name}', value=pdfextractdata['metadata'].get('doi', ''))
+                                    page = st.text_input('Pages', key=f'page_{uploaded_file.name}', value=pdfextractdata['metadata'].get('page', ''))
                                 submitted = st.form_submit_button('Save Metadata')
-
                                 if submitted:
-                                    add_metadata(file_path, title, author, year, publisher, DOI,page,volume,issue)
+                                    add_metadata(file_path, title, author, year, journal,publisher, DOI,page,volume,issue)
                                     st.success('Metadata saved!')
 
                     with tab2:
@@ -338,6 +367,7 @@ def main():
                                         "title": corresponding_metadata["title"],
                                         "author": corresponding_metadata["author"],
                                         "year": corresponding_metadata["year"],
+                                        "journal": corresponding_metadata["journal"],
                                         "publisher": corresponding_metadata["publisher"],
                                         "DOI": corresponding_metadata["DOI"],
                                         "page": corresponding_metadata["page"],