Update inference.py

inference.py

@@ -6,17 +6,15 @@ from peft import PeftModel, PeftConfig
 import torch
 
 orig_checkpoint = 'google/gemma-2b'
-checkpoint = 'checkpoint-4000'
+checkpoint = '.'
 HF_TOKEN = ''
 PROMPT = 'Salut, ca sa imi schimb buletinul pot sa'
 
-seq_len = 2048
+seq_len = 256
 
 # load original model first
 tokenizer = AutoTokenizer.from_pretrained(orig_checkpoint, token=HF_TOKEN)
-
-config = PeftConfig.from_pretrained(checkpoint)
-model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(orig_checkpoint, token=HF_TOKEN)
 
 # then merge trained QLoRA weights
 model = PeftModel.from_pretrained(model, checkpoint)
@@ -28,4 +26,4 @@ model = model.cuda()
 inputs = tokenizer.encode(PROMPT, return_tensors="pt").cuda()
 outputs = model.generate(inputs, max_new_tokens=seq_len)
 
-print(tokenizer.decode(outputs[0]))
+print(tokenizer.decode(outputs[0]))