Update README.md

README.md

@@ -67,11 +67,74 @@ Users should validate the simplified text with healthcare professionals to ensur
 
 Use the code below to get started with the model.
 
-'''python
+```python
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig
+)
+
+from peft import PeftConfig
+
+MODEL = "9rofe/Wernicke-AI3"
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+config = PeftConfig.from_pretrained(MODEL)
+model = AutoModelForCausalLM.from_pretrained(
+    config.base_model_name_or_path,
+    return_dict=True,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+tokenizer=AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+tokenizer.pad_token = tokenizer.eos_token
+
+model = PeftModel.from_pretrained(model, MODEL)
+
+generation_config = model.generation_config
+generation_config.max_new_tokens = 500 # MODIFY
+generation_config.temperature = 0.7
+generation_config.top_p = 0.7
+generation_config.num_return_sequences = 1
+generation_config.pad_token_id = tokenizer.eos_token_id
+generation_config.eos_token_id = tokenizer.eos_token_id
+
+%%time
+device = "cuda:0"
+
+prompt = """
+<user>: Convert this text to reading level 6: {TEXT}
+<assistant>:
+""".strip()
+
+encoding = tokenizer(prompt, return_tensors="pt").to(device)
+with torch.inference_mode():
+  outputs = model.generate(
+      input_ids = encoding.input_ids,
+      attention_mask = encoding.attention_mask,
+      generation_config = generation_config
+  )
+
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+Utilize this prompt:
+
+```python
 prompt = """
-<user>: Convert this text to reading level 6: What are the complications? What are the complications? Nearly 88,000 people die from alcohol-related causes each year in the US, making alcohol use the 4th leading preventable cause of death. Drinking at an increased risk level raises your chance of:  Accidents, injuries, and aggression. Drinking too much increases your risk for every type of injury and violence. Alcohol is a factor in about 60% of fatal burn injuries, drownings, and murders; 50% of severe trauma injuries and sexual assaults; and 40% of fatal crashes and falls.  Physical health problems. Heavy drinkers have a greater chance of liver and heart disease, stroke, digestive problems, and some types of cancer. Theyre also more likely to have problems with sexual function and premature aging.  Emotional and cognitive problems. People who drink too much are more prone to anxiety and depression. They may have trouble sleeping, remembering things, and solving problems.  Problems with relationships, work, and studies. Heavy drinking can interfere with your interactions and performance in every area of your life.  Birth defects. Drinking during pregnancy can cause brain damage and deformities in the baby. Since scientists dont know whether any amount of alcohol is safe for a developing baby, women who are pregnant or trying to become pregnant should not drink.  Alcoholism and alcohol use disorders. Drinking at an increased risk level raises your chance of developing an alcohol use disorder. How is it treated? Your medical provider may refer you to an addiction specialist for consultation if there is a perceived risk. If your risk is moderate to high, you may need to have treatment \(at an outpatient or in an acute hospital setting\) to help you manage withdrawal symptoms and reduce your risks associated with this disorder. Your medical provider may also prescribe medicines to help you to manage your symptoms. 
+<user>: Convert this text to reading level 6: {text}
 <assistant>:
-""".strip()'''
+""".strip()
+```
 
 ## Training Details
 
@@ -151,7 +214,7 @@ Carbon emissions can be estimated using the [Machine Learning Impact calculator]
 - **Hardware Type:** GPU (NVIDIA A100)
 - **Hours used:** 120 hours
 - **Cloud Provider:** AWS
-- **Compute Region:** US East (N. Virginia)
+- **Compute Region:** US West (Utah)
 - **Carbon Emitted:** 500 kg CO2eq
 
 ## Technical Specifications [optional]