Update README.md

README.md

@@ -53,24 +53,50 @@ Model evaluation was based on qualitative assessment of generated text relevance
 Here is how to load and use the model:
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-model_name = "vanessasml/cyber-risk-llama-3-8b"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Example of how to use the model:
-prompt = """Question: What are the cyber threads present in the article?
+model_id = "vanessasml/cyber-risk-llama-3-8b"
+
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model_id,
+    model_kwargs={"torch_dtype": torch.bfloat16},
+    device="cuda",
+)
+## Define your user prompt
+example_prompt_1=""" Question: What are the cyber threats present in the article?Explain why.\n
 Article: More than one million Brits over the age of 45 have fallen victim to some form of email-related fraud, \
 as the internet supersedes the telephone as the favored channel for scammers, according to Aviva. \
 The insurer polled over 1000 adults over the age of 45 in the latest update to its long-running Real Retirement Report. \
-Further, 6% said they had actually fallen victim to such an online attack, amounting to around 1.2 million adults. \
-Some 22% more people it surveyed had been targeted by ...	
+Further, 6% said they had actually fallen victim to such an online attack, amounting to around 1.2 million adults. 	
 """
-pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
-# To generate text:
-result = pipe(prompt)
-print(result[0]['generated_text'])
+example_prompt_2 = "What are the main 5 cyber classes from the NIST cyber framework?" 
+
+messages = [
+    {"role": "system", "content": "You are an IT supervisor from a supervisory institution."},
+    {"role": "user", "content": example_prompt_2},
+]
+
+prompt = pipeline.tokenizer.apply_chat_template(
+        messages, 
+        tokenize=False, 
+        add_generation_prompt=True
+)
+
+terminators = [
+    pipeline.tokenizer.eos_token_id,
+    pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
+]
+
+outputs = pipeline(
+    prompt,
+    max_new_tokens=500,
+    eos_token_id=terminators,
+    do_sample=True,
+    temperature=0.1,
+    top_p=0.9,
+)
+print(outputs[0]["generated_text"][len(prompt):])
+
+## Example output
 ```
 
 ## Limitations and Bias
@@ -81,7 +107,7 @@ The model, while robust in cybersecurity contexts, may not generalize well to un
 If you use this model, please cite it as follows:
 
 ```bibtex
-@misc{cyber-risk-llama-3-8b-sft-lora-4bit-float16,
+@misc{cyber-risk-llama-3-8b,
   author = {Vanessa Lopes},
   title = {Cyber-risk-llama-3-8B Model},
   year = {2024},