Update README.md

README.md

@@ -29,8 +29,8 @@ There was great political strife in the air, and tensions were high. People rush
         model = GPTJForCausalLM.from_pretrained(pretrain_name, load_in_8bit=load_in_8bit, device_map='auto', torch_dtype=torch.float16)
         model = PeftModel.from_pretrained(model,lora_weights,torch_dtype=torch.float16,device_map={'':0})
         model = torch.compile(model)
-
-        GenerationConfig(
+        tokenizer = AutoTokenizer.from_pretrained('pygmalion-6b') #The orginal pretrained
+        gen_config=GenerationConfig(
                 temperature=0.1,
                 top_p=0.75,
                 top_k=40,
@@ -38,11 +38,11 @@ There was great political strife in the air, and tensions were high. People rush
             )
 
         text = '[User]: What's the best food in Hanoi?''
-        input_ids = st.session_state.tokenizer(text, return_tensors='pt')['input_ids'].to('cuda')
+        input_ids = tokenizer(text, return_tensors='pt')['input_ids'].to('cuda')
         with torch.no_grad():
-            output = st.session_state['model'].generate(input_ids=input_ids, generation_config=st.session_state.gen_config,return_dict_in_generate=True, output_scores=True,max_new_tokens=256)
+            output = model.generate(input_ids=input_ids, generation_config=gen_config,return_dict_in_generate=True, output_scores=True,max_new_tokens=256)
             s = output.sequences[0]
-            output = st.session_state.tokenizer.decode(s)
+            output = tokenizer.decode(s)
             print('Raw:',output)
         
 ```