Updates

setup.py

@@ -12,7 +12,7 @@ with open('requirements.txt') as f:
 setup(
     name='t5s',
     packages=find_packages(include=['t5s*']),
-    version='2.0.3',
+    version='2.0.4',
     description="T5 Summarisation Using Pytorch Lightning",
     license='MIT License',
     classifiers=[

src/data/process_data.py

@@ -9,7 +9,7 @@ def process_data(split="train"):
 
     df = pd.read_csv("data/raw/{}.csv".format(split))
     df.columns = ["Unnamed: 0", "input_text", "output_text"]
-    df = df.sample(frac=float(params["split"]), replace=True, random_state=1)
+    df = df.sample(frac=params["split"], replace=True, random_state=1)
     df.to_csv("data/processed/{}.csv".format(split))
 
 

src/models/predict_model.py

@@ -11,6 +11,6 @@ def predict_model(text: str):
         params = yaml.safe_load(f)
 
     model = Summarization()
-    model.load_model(model_type=params["model_type"], model_dir=params["model_dir"])
+    model.load_model(model_type=params["model_type"], model_dir="gagan3012/summarsiation")
     pre_summary = model.predict(text)
     return pre_summary

t5s/cli.py

@@ -22,31 +22,32 @@ parser_start.add_argument(
     "-d",
     "--dataset",
     default="cnn_dailymail",
-    help="Enter the name of the dataset to be used",
+    help="Enter the name of the dataset to be used",type=str
 )
 
-parser_start.add_argument("-s", "--split", default=0.001, help="Enter the split required")
+parser_start.add_argument("-s", "--split", default=0.001, help="Enter the split required",type=float)
 
 parser_start.add_argument(
     "-n", "--name", default="summarsiation", help="Enter the name of the model"
 )
 parser_start.add_argument(
-    "-mt", "--model_type", default="t5", help="Enter the model type"
+    "-mt", "--model_type", default="t5", help="Enter the model type",type=str
 )
 parser_start.add_argument(
     "-m",
     "--model_name",
     default="t5-base",
     help="Enter the model to be used eg t5-base",
+    type=str
 )
 parser_start.add_argument(
-    "-e", "--epochs", default=5, help="Enter the number of epochs"
+    "-e", "--epochs", default=5, help="Enter the number of epochs", type=int
 )
 parser_start.add_argument(
-    "-lr", "--learning-rate", default=0.0001, help="Enter the number of epochs"
+    "-lr", "--learning-rate", default=0.0001, help="Enter the number of epochs", type=float
 )
 parser_start.add_argument(
-    "-b", "--batch-size", default=2, help="Enter the number of batches"
+    "-b", "--batch-size", default=2, help="Enter the number of batches", type=int
 )
 
 parser_dirs = command_subparser.add_parser(