Terry Zhang
commited on
Commit
·
7688055
1
Parent(s):
94cab49
add prints
Browse files- tasks/text.py +6 -0
tasks/text.py
CHANGED
@@ -33,6 +33,7 @@ def baseline_model(dataset_length: int):
|
|
33 |
return predictions
|
34 |
|
35 |
def tree_classifier(test_dataset: dict, model: str):
|
|
|
36 |
|
37 |
texts = test_dataset["quote"]
|
38 |
|
@@ -47,6 +48,8 @@ def tree_classifier(test_dataset: dict, model: str):
|
|
47 |
'xgboost.sklearn.XGBClassifier'])
|
48 |
|
49 |
predictions = model.predict(texts)
|
|
|
|
|
50 |
|
51 |
return predictions
|
52 |
|
@@ -71,6 +74,7 @@ class TextDataset(Dataset):
|
|
71 |
|
72 |
|
73 |
def bert_classifier(test_dataset: dict, model: str):
|
|
|
74 |
texts = test_dataset["quote"]
|
75 |
|
76 |
model_repo = f"theterryzhang/frugal_ai_{model}"
|
@@ -97,6 +101,8 @@ def bert_classifier(test_dataset: dict, model: str):
|
|
97 |
p = torch.argmax(outputs.logits, dim=1)
|
98 |
predictions = np.append(predictions, p.cpu().numpy())
|
99 |
|
|
|
|
|
100 |
return predictions
|
101 |
|
102 |
|
|
|
33 |
return predictions
|
34 |
|
35 |
def tree_classifier(test_dataset: dict, model: str):
|
36 |
+
print("Starting tree model run")
|
37 |
|
38 |
texts = test_dataset["quote"]
|
39 |
|
|
|
48 |
'xgboost.sklearn.XGBClassifier'])
|
49 |
|
50 |
predictions = model.predict(texts)
|
51 |
+
print("Finished tree model run")
|
52 |
+
|
53 |
|
54 |
return predictions
|
55 |
|
|
|
74 |
|
75 |
|
76 |
def bert_classifier(test_dataset: dict, model: str):
|
77 |
+
print("Starting BERT model run")
|
78 |
texts = test_dataset["quote"]
|
79 |
|
80 |
model_repo = f"theterryzhang/frugal_ai_{model}"
|
|
|
101 |
p = torch.argmax(outputs.logits, dim=1)
|
102 |
predictions = np.append(predictions, p.cpu().numpy())
|
103 |
|
104 |
+
print("Finished BERT model run")
|
105 |
+
|
106 |
return predictions
|
107 |
|
108 |
|