Update README.md

README.md

@@ -25,6 +25,77 @@ Iris provides efficient and accurate translation and can be used in a variety of
 * **base mode** : mistralai/Mistral-7B-v0.1
 * **dataset** : translation_v3_346k
 
+## usage
+```
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from utils.simple_bleu import simple_score
+import torch
+
+repo = "davidkim205/iris-7b"
+model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype=torch.bfloat16, device_map='auto')
+tokenizer = AutoTokenizer.from_pretrained(repo)
+
+def generate(prompt):
+    encoding = tokenizer(
+        prompt,
+        return_tensors='pt',
+        return_token_type_ids=False
+    ).to("cuda")
+    gen_tokens = model.generate(
+        **encoding,
+        max_new_tokens=2048,
+        temperature=1.0,
+        num_beams=5,
+    )
+    prompt_end_size = encoding.input_ids.shape[1]
+    result = tokenizer.decode(gen_tokens[0, prompt_end_size:])
+
+    return result
+
+
+def translate_ko2en(text):
+    prompt = f"[INST] 다음 문장을 영어로 번역하세요.{text} [/INST]"
+    return generate(prompt)
+
+
+def translate_en2ko(text):
+    prompt = f"[INST] 다음 문장을 한글로 번역하세요.{text} [/INST]"
+    return generate(prompt)
+
+
+def main():
+    while True:
+        text = input('>')
+        en_text = translate_ko2en(text)
+        ko_text = translate_en2ko(en_text)
+        print('en_text', en_text)
+        print('ko_text', ko_text)
+        print('score', simple_score(text, ko_text))
+
+if __name__ == "__main__":
+    main()
+```
+output
+```
+$ python iris_test.py 
+Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████| 3/3 [00:02<00:00,  1.05it/s]
+>Iris is a model for Korean-English sentence translation based on deep learning.
+Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.
+Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.
+en_text Iris is a model for Korean-English sentence translation based on deep learning.</s>
+ko_text 아이리스는 딥러닝을 기반으로 한 한-영어 문장 번역을 위한 모델이다.</s>
+
+```
+## template
+### ko -> en
+```
+[INST] 다음 문장을 영어로 번역하세요.{text} [/INST]
+```
+### en -> ko
+```
+"[INST] 다음 문장을 한글로 번역하세요.{text} [/INST]"
+```
+
 ## dataset info : translation_v3_346k
 
 | dataset name                  | ratio  | size          |