Update README.md

README.md

@@ -68,8 +68,10 @@ A model finetuned with the [NoticIA Dataset](https://huggingface.co/datasets/Ike
 
 # Usage example:
 ```python
-from datasets import load_dataset
-from transformers import pipeline
+import torch # pip install torch
+from datasets import load_dataset # pip install datasets
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig # pip install transformers
+
 
 def prompt(
     headline: str,
@@ -103,21 +105,37 @@ def prompt(
         f"{body}\n"
     )
 
+dataset = load_dataset("Iker/NoticIA")
+example = dataset["test"][0]
 
+prompt = prompt(headline=example["web_headline"], body=example["web_text"])
 
+tokenizer = AutoTokenizer.from_pretrained("Iker/ClickbaitFighter-2B")
+model = AutoModelForCausalLM.from_pretrained(
+    "Iker/ClickbaitFighter-2B", torch_dtype=torch.bfloat16, device_map="auto"
+)
 
+formatted_prompt = tokenizer.apply_chat_template(
+    [{"role": "user", "content": prompt}],
+    tokenize=False,
+    add_generation_prompt=True,
+)
 
-dataset = load_dataset("Iker/NoticIA")
-
-example = dataset["test"][0]
-
-model_input = prompt(headline=example["web_headline"],body=example["web_text"])
+model_inputs = tokenizer(
+    [formatted_prompt], return_tensors="pt", add_special_tokens=False
+)
 
-pipe = pipeline("text-generation", model="Iker/ClickbaitFighter-2B",device_map="auto")
+model_output = model.generate(**model_inputs.to(model.device), generation_config=GenerationConfig(
+  max_new_tokens=32,
+  min_new_tokens=1,
+  do_sample=False,
+  num_beams=1,
+  use_cache=True
+))
 
-summary = pipe(model_input)
+summary = tokenizer.batch_decode(model_output,skip_special_tokens=True)[0]
 
-print(summary)
+print(summary.strip().split("\n")[-1]) # Get only the summary, without the prompt. 
 ```
 
 # Evaluation Results