from phenotype_inference import PhenotypeClassifier, print_prediction

MODEL_PATH = "phenotype_tinytransformer_v1_temperature_scaled.pt"

classifier = PhenotypeClassifier(MODEL_PATH)

features = {
    "Gram Stain": "Negative",
    "Shape": "Rods",
    "Catalase": "Positive",
    "Oxidase": "Positive",
    "Motility": "Positive",
    "Indole": "Negative",
    "Citrate": "Positive",
    "Urease": "Negative",
    "Growth Temperature": "20//37",
    "Media Grown On": "Blood Agar; MacConkey Agar",
}

result = classifier.predict(features, top_k=10)

print_prediction(result)

print()
print("Raw output:")
print(result)