Update README.md

README.md

@@ -46,7 +46,7 @@ You can run it on [Colab](https://colab.research.google.com/drive/1FgdwkkPcLzn_x
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, LlamaTokenizer, GenerationConfig, pipeline
 from peft import PeftModel, PeftMixedModel
 import torch
-import pprint
+
 
 model_name = "0xtaipoian/open-lilm"
 
@@ -59,13 +59,12 @@ bnb_config = BitsAndBytesConfig(
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
   model_name,
+  revision="main", #qlora-merged (qLoRA finetuned for 3 epochs) or main (full parameter finetune for 1 epoch)
   torch_dtype=torch.bfloat16,
-  device_map='auto',
   trust_remote_code=True,
   quantization_config=bnb_config,
 )
 
-
 def chat(messages, temperature=0.9, max_new_tokens=200):
     input_ids = tokenizer.apply_chat_template(conversation=messages, tokenize=True, add_generation_prompt=True, return_tensors='pt').to('cuda:0')
     output_ids = model.generate(input_ids, max_new_tokens=max_new_tokens, temperature=temperature, do_sample=True)
@@ -83,13 +82,7 @@ messages = [
 
              "content":
 """
-密陽44人輪姦案」受害女隔20年現身：時間停在2004，不記...
-
-南韓密陽2004年曾發生一起44人輪姦一位14歲少女案件，44名加害者當時幾乎都未成年，
-法官輕判引起社會聲浪抗議。案件如今過了20年，當年只有14歲的少女與妹妹相隔20年現
-身，近日接受節目《想知道真相》訪問，坦言沒有藥物無法生活，時間就此停在了2004年
-走不出來，「我的夢想是什麼，我不記得了。」只希望家人們從惡夢裡逃脫，過著幸福日
-子。
+密陽44人輪姦案」受害女隔20年現身：時間停在2004，不記得
 """}]
 
 result = chat(messages, max_new_tokens=200, temperature=1)