castorini
/

rankllama-v1-7b-lora-doc

Model card Files Files and versions Community

MrLight commited on Oct 24, 2023

Commit

5e9e894

•

1 Parent(s): 2435db0

Update README.md

Files changed (1) hide show

README.md +1 -1

README.md CHANGED Viewed

@@ -24,7 +24,7 @@ from peft import PeftModel, PeftConfig
 def get_model(peft_model_name):
     config = PeftConfig.from_pretrained(peft_model_name)
-    base_model = AutoModelForSequenceClassification.from_pretrained(config.base_model_name_or_path)
     model = PeftModel.from_pretrained(base_model, peft_model_name)
     model = model.merge_and_unload()
     model.eval()

 def get_model(peft_model_name):
     config = PeftConfig.from_pretrained(peft_model_name)
+    base_model = AutoModelForSequenceClassification.from_pretrained(config.base_model_name_or_path, num_labels=1)
     model = PeftModel.from_pretrained(base_model, peft_model_name)
     model = model.merge_and_unload()
     model.eval()