Change sktm handling of CSV files

jilljenn · Jan 10, 2024 · bd32da5 · bd32da5
1 parent 446e851
commit bd32da5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -89,7 +89,7 @@ Usage:
 
     mkdir data/assistments09
     wget https://jiji.cat/weasel2018/data.csv -P data/assistments09
-    python sktm.py --dataset assistments09 --model (irt|pfa|sktm)  # Choose which model
+    python sktm.py data/assistments09/data.csv --model (irt|pfa|sktm)  # Choose which model
 
 For factorization machines, replace `LogisticRegression` with `from fm import FMClassifier`.
 

diff --git a/sktm.py b/sktm.py
@@ -10,12 +10,13 @@
 
 
 parser = argparse.ArgumentParser(description='Run simple KTM')
-parser.add_argument('--dataset', type=str, nargs='?', default='dummy')
+parser.add_argument(
+    'csv_file', type=str, nargs='?', default='data/dummy/data.csv')
 parser.add_argument('--model', type=str, nargs='?', default='iswf')
 options = parser.parse_args()
 
 
-df = pd.read_csv(f'data/{options.dataset}/data.csv')
+df = pd.read_csv(options.csv_file)
 pipe = Pipeline([
     ('onehot', OneHotEncoder(handle_unknown='ignore')),
     ('lr', LogisticRegression(solver='liblinear'))