init🎉:

test_prediction.py

@@ -0,0 +1,24 @@
+import json
+from realign.eval_acc import LM
+
+
+device = 'cuda:1'
+model_path = './ckpts/stylish'
+data_path = './data/nyt10/nyt10_test.txt'
+
+prompt_template = (
+    "Below is an instruction that describes a task, paired with an input that provides further context. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\nGiven a piece of text, please find out the person-nationality relation in it. Tell me who is the person and which is the nationality. The answer should be in json format.\n\n### Input:\n{text}\n\n### Response:"
+)
+
+
+data = [json.loads(line) for line in open(data_path, 'r').readlines()]
+data = [dp for dp in data if '/people/person/nationality' in dp['relation']]
+
+lm = LM(device, model_path)
+
+for i in range(len(data[:10])):
+    prompt = prompt_template.format(text=data[i]['text'])
+    response = lm.chat(prompt)
+    print(response)