Update README.md

README.md

@@ -31,38 +31,21 @@ This model has been fine-tuned using 4-bit QLORA, based on [Llama-3-8B from Meta
 
 The model can be loaded with HuggingFace's Transformers library:
 ``` python
-import transformers
-import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_id = "duke-nlp/Prob-Gen-8B"
+model_id = "DukeNLP/Prob-Gen-8B"
 
-model = AutoModelForCausalLM.from_pretrained(
-	model_id,
-	device_map="auto",
-	torch_dtype=torch.bfloat16
-)
+model = AutoModelForCausalLM.from_pretrained(model_id,device_map="auto", trust_remote_code=True)
 
-tokenizer = AutoTokenizer.from_pretrained(
-	model_id, 
-	use_fast=False,
-	legacy=False
-)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-model_input = tokenizer(
-	"""Please generate a math problem and 2 to 4 options for 8th graders with the following requirements:
-Problem context: <specified-context>
-Tested knowledge: <specified-knowledge>""", 
-	return_tensors="pt",
-).to("cuda")
-
-model_output = model.generate(
-	model_input['input_ids'], 
-	max_new_tokens=256, 
-	do_sample=True,
-	...
-)
-
-tokenizer.batch_decode(model_output)
+prompt = "Please generate a math problem and 2 to 4 options for 8th graders with the following requirements:\nProblem context: <specified-context>\nTested knowledge: <specified-knowledge>"
+
+model_input = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+model_output = model.generate(model_input['input_ids'], max_new_tokens=256)
+
+print(tokenizer.batch_decode(model_output))
 ```
 
 <!-- ## Bias, Risks, and Limitations