-
Notifications
You must be signed in to change notification settings - Fork 269
/
test_ci.py
35 lines (28 loc) · 1.07 KB
/
test_ci.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
#!/usr/bin/env python3
"""
Tests of ktrain text classification flows
"""
import testenv
import numpy as np
from unittest import TestCase, main, skip
from ktrain import tabular
def adult_census():
import pandas as pd
df = pd.read_csv('tabular_data/adults.csv')
df = df.rename(columns=lambda x: x.strip())
df = df.applymap(lambda x: x.strip() if isinstance(x, str) else x)
filter_set = 'Doctorate'
df['treatment'] = df['education'].apply(lambda x: 1 if x in filter_set else 0)
return df
class TestCausalInference(TestCase):
def test_ci(self):
df = adult_census()
cm = tabular.causalinference.causal_inference_model(df, metalearner_type='t-learner',
treatment_col='treatment',
outcome_col='class',
ignore_cols=['fnlwgt', 'education','education-num']).fit()
ate = cm.estimate_ate()
self.assertGreater(ate['ate'], 0.20)
self.assertLess(ate['ate'], 0.21)
if __name__ == "__main__":
main()