Update app.py

app.py

@@ -12,6 +12,12 @@ import torch
 import shutil
 import os
 import uuid
+import json
+
+
+from google.oauth2 import service_account
+import gspread
+from google.oauth2.service_account import Credentials
 
 
 css = """
@@ -21,7 +27,6 @@ css = """
 """
 st.markdown(css, unsafe_allow_html=True)
 
-
 def main():
     st.title("💡 LoraHub")
     st.markdown("Low-rank adaptations (LoRA) are techniques for fine-tuning large language models on new tasks. We propose LoraHub, a framework that allows composing multiple LoRA modules trained on different tasks. The goal is to achieve good performance on unseen tasks using just a few examples, without needing extra parameters or training. And we want to build a marketplace where users can share their trained LoRA modules, thereby facilitating the application of these modules to new tasks.")
@@ -105,12 +110,28 @@ Infer the date from context.  Q: Today is the second day of the third month of 1
                                                   txt_input, txt_output, max_inference_step=max_step)
                 
             st.success("Lorahub learning finished! You got the following recommendation:")
+
             df = {
                 "modules": [LORA_HUB_NAMES[i] for i in st.session_state["select_names"]],
                 "weights": recommendation.value,
             }
+
+
+
+            def share():
+                credentials = service_account.Credentials.from_service_account_info(
+                json.loads(st.secrets["gcp_service_account"]),
+                scopes=[
+                    "https://www.googleapis.com/auth/spreadsheets",
+                ]
+                )
+                gsheet_url = st.secrets["private_gsheets_url"]
+                gc = gspread.authorize(credentials)
+                sh = gc.open_by_url(gsheet_url)
+
+                ws = sh.sheet1
+                ws.insert_rows([[LORA_HUB_NAMES[i] for i in st.session_state["select_names"]],recommendation.value.tolist(),[]])
             st.table(df)
-            
             random_id = uuid.uuid4().hex
             os.makedirs(f"lora/{random_id}")
             # copy config file
@@ -126,7 +147,16 @@ Infer the date from context.  Q: Today is the second day of the third month of 1
                     file_name=f"lora_{random_id}.zip",
                     mime="application/zip"
                 )
-            st.warning("The page will be refreshed once you click the download button.")
+            with open(f"lora_{random_id}.zip", "rb") as fp:
+                btn = st.download_button(
+                    label="📥 Download the final LoRA Module and share your results",
+                    data=fp,
+                    file_name=f"lora_{random_id}.zip",
+                    mime="application/zip",
+                    on_click=share
+                )
+            st.button("📥 Share your results",on_click=share)
+            st.warning("The page will be refreshed once you click the download button. Share results may cost 1-2 mins.")