Add option --no-used-evidence
This commit is contained in:
parent
d96d1fc96e
commit
e6dab8395f
|
@ -32,7 +32,7 @@ def add_message(msg):
|
|||
msg.get_payload(decode=True)
|
||||
.decode(charset, errors='replace'))
|
||||
|
||||
def extract_features(msgid, verbose):
|
||||
def extract_features(msgid, verbose, used_evidence):
|
||||
db = psycopg2.connect("dbname=bayes")
|
||||
csr = db.cursor(cursor_factory=ppe.DictCursor)
|
||||
csr.execute(
|
||||
|
@ -87,6 +87,7 @@ def extract_features(msgid, verbose):
|
|||
for i in interesting_evidence:
|
||||
if verbose:
|
||||
print("#", i["spam_prob"], i["type"], i["length"], i["feature"], sep="\t")
|
||||
if used_evidence:
|
||||
csr.execute(
|
||||
"insert into used_evidence(message, spam_prob, type, length, feature) values(%s, %s, %s, %s, %s)",
|
||||
(msg_pk, i["spam_prob"], i["type"], i["length"], i["feature"]))
|
||||
|
@ -100,6 +101,7 @@ def extract_features(msgid, verbose):
|
|||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument('--verbose', action='store_true')
|
||||
ap.add_argument('--no-used-evidence', action='store_false', dest='used_evidence')
|
||||
ap.add_argument('file', nargs='?')
|
||||
args = ap.parse_args()
|
||||
if args.file:
|
||||
|
@ -109,7 +111,7 @@ def main():
|
|||
parser = email.parser.BytesParser(policy=email.policy.default)
|
||||
msg = parser.parse(fh)
|
||||
add_message(msg)
|
||||
p = extract_features(msg["Message-Id"], args.verbose)
|
||||
p = extract_features(msg["Message-Id"], args.verbose, args.used_evidence)
|
||||
print(p, "spam" if p > 0.5 else "ham")
|
||||
|
||||
main()
|
||||
|
|
Loading…
Reference in New Issue