Update app.py

app.py

@@ -6,7 +6,7 @@ from datasets import load_dataset
 from trl import SFTTrainer
 
 ref_model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-70m-deduped-v0", torch_dtype=torch.bfloat16)
-peft_model_id = "w601sxs/pythia-70m-instruct-orca-chkpt-64000"
+peft_model_id = "w601sxs/pythia-70m-instruct-orca-chkpt-1245000"
 
 config = PeftConfig.from_pretrained(peft_model_id)
 model = PeftModel.from_pretrained(ref_model, peft_model_id)
@@ -17,8 +17,8 @@ model.eval()
 def predict(text):
     inputs = tokenizer(text, return_tensors="pt")
     with torch.no_grad():
-        outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=10)
-        out_text = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]
+        outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=128)
+        out_text = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0].split("answer:")[-1]
     
     return out_text.split(text)[-1]