Terry Zhang commited on
Commit
7688055
·
1 Parent(s): 94cab49

add prints

Browse files
Files changed (1) hide show
  1. 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