Update compute_model_property.py

compute_model_property.py

@@ -16,6 +16,8 @@ import numpy as np
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
 from evaluate import load
 
+from fastapi import FastAPI
+app = FastAPI()
 
 # 1. record each file name included
 # 1.1 read different file formats depending on parameters (i.e., filetype)
@@ -267,8 +269,9 @@ def compute_model_card_evaluation_results(tokenizer, model_checkpoint, raw_datas
     result = trainer.evaluate()
     return result
 
-
-if __name__ == "__main__":
+#if __name__ == "__main__":
+@app.get("/")
+def return_output():
     
     in_container = True
     if len(sys.argv) > 1:
@@ -284,16 +287,18 @@ if __name__ == "__main__":
 
     print(model_checkpoint, dataset_name, metric)
     
-
     model_checkpoint = model_checkpoint
     raw_datasets = load_dataset(dataset_name, "mrpc")
     metric = load("glue", "mrpc")
     tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
     output = compute_model_card_evaluation_results(tokenizer, model_checkpoint, raw_datasets, metric)
-    print(json.dumps(output))
-
-    if in_container:
-        with open("/tmp/outputs/computation_result.json", "w") as f:
-            json.dump(output, f, indent=4, sort_keys=True)
-    else:
-        print(json.dumps(output, indent=4, sort_keys=True))
+    
+    #print(json.dumps(output))
+       
+    #if in_container:
+    #    with open("/tmp/outputs/computation_result.json", "w") as f:
+    #        json.dump(output, f, indent=4, sort_keys=True)
+    #else:
+    #    print(json.dumps(output, indent=4, sort_keys=True))
+
+    return output