updates

params.yml

@@ -1,3 +1,4 @@
+name: summarsiation
 data: cnn_dailymail
 batch_size: 2
 num_workers: 2
@@ -8,6 +9,8 @@ epochs: 5
 source_dir: src
 model_dir: models
 metric: rouge
-split: 0.02
+split: 0.001
 use_gpu: True
-visualise: True
+visualise: True
+hf_username: gagan3012
+upload_to_hf: True

src/__init__.py

@@ -1,8 +1,12 @@
-from src.models.evaluate_model import evaluate_model
-from src.models.predict_model import predict_model
-from src.models.train_model import train_model
-from src.models.model import Summarization
-from src.data.make_dataset import make_dataset
-from src.data.process_data import process_data
-from src.visualization.visualize import visualize
-import os, sys; sys.path.append(os.path.dirname(os.path.realpath(__file__)))
+import os  # noqa: F401
+import sys  # noqa: F401
+
+from src.data.make_dataset import make_dataset  # noqa: F401
+from src.data.process_data import process_data  # noqa: F401
+from src.models.evaluate_model import evaluate_model  # noqa: F401
+from src.models.model import Summarization  # noqa: F401
+from src.models.predict_model import predict_model  # noqa: F401
+from src.models.train_model import train_model  # noqa: F401
+from src.visualization.visualize import visualize  # noqa: F401
+
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))  # noqa: F401

src/models/__init__.py

@@ -1,4 +1,4 @@
-from .model import Summarization
-from .train_model import train_model
-from .predict_model import predict_model
-from .evaluate_model import evaluate_model
+from .model import Summarization  # noqa: F401
+from .train_model import train_model  # noqa: F401
+from .predict_model import predict_model  # noqa: F401
+from .evaluate_model import evaluate_model  # noqa: F401

src/models/evaluate_model.py

@@ -1,4 +1,3 @@
-import dagshub
 import yaml
 
 from model import Summarization

src/models/model.py

@@ -1,5 +1,10 @@
+import shutil
+from getpass import getpass
+from pathlib import Path
+
 import torch
 import pandas as pd
+from huggingface_hub import HfApi, Repository
 from transformers import (
     AdamW,
     T5ForConditionalGeneration,
@@ -547,3 +552,31 @@ class Summarization:
             "rougeLsum High F1": results["rougeLsum"].high.fmeasure,
         }
         return output
+
+    def upload(self, hf_username, model_name):
+        hf_password = getpass("Enter your HuggingFace password")
+        if Path("./models").exists():
+            shutil.rmtree("./models")
+        token = HfApi().login(username=hf_username, password=hf_password)
+        del hf_password
+        model_url = HfApi().create_repo(token=token, name=model_name, exist_ok=True)
+        model_repo = Repository(
+            "./model",
+            clone_from=model_url,
+            use_auth_token=token,
+            git_email=f"{hf_username}@users.noreply.huggingface.co",
+            git_user=hf_username,
+        )
+
+        readme_txt = f"""
+            ---
+            Summarisation model {model_name}
+            """.strip()
+
+        (Path(model_repo.local_dir) / "README.md").write_text(readme_txt)
+        self.save_model()
+        commit_url = model_repo.push_to_hub()
+
+        print("Check out your model at:")
+        print(commit_url)
+        print(f"https://huggingface.co/{hf_username}/{model_name}")

src/models/train_model.py

@@ -43,6 +43,9 @@ def train_model():
     with open("reports/training_metrics.txt", "w") as fp:
         json.dump(data, fp)
 
+    if params["upload_to_hf"]:
+        model.upload(hf_username=params["hf_username"], model_name=params["name"])
+
 
 if __name__ == "__main__":
     train_model()

src/visualization/visualize.py

@@ -5,12 +5,13 @@ from models import predict_model
 
 
 def visualize():
-    st.write('# Summarization  UI')
+    st.write("# Summarization  UI")
     st.markdown(
-        '''
+        """
         *For additional questions and inquiries, please contact **Gagan Bhatia** via [LinkedIn](
         https://www.linkedin.com/in/gbhatia30/) or [Github](https://github.com/gagan3012).*
-        ''')
+        """
+    )
 
     text = st.text_area("Enter text here")
     if st.button("Generate Summary"):
@@ -27,5 +28,5 @@ if __name__ == "__main__":
     with open("params.yml") as f:
         params = yaml.safe_load(f)
 
-    if params['visualise']:
+    if params["visualise"]:
         visualize()