Update README.md

README.md

@@ -9,6 +9,7 @@ metrics:
   - rouge
 library_name: transformers
 pipeline_tag: text-generation
+base_model: google/gemma-2b-it
 tags:
   - clickbait
   - noticia
@@ -48,7 +49,6 @@ widget:
           la señora ha encantado a los usuarios de la red. Es por eso que el
           relato ha acumulado más de 1.000 me gusta y cerca de 100 retuits,
           además de una multitud de comentarios.\\n"
-
 ---
 
 <table>
@@ -69,11 +69,19 @@ If you are looking for a larger model, with better performance, check out [Click
 
 
 # Usage example:
+
+## Summarize a web article
 ```python
 import torch # pip install torch
-from datasets import load_dataset # pip install datasets
+from newspaper import Article #pip3 install newspaper3k
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig # pip install transformers
 
+article_url ="https://www.huffingtonpost.es/virales/le-compra-abrigo-abuela-97nos-reaccion-fantasia.html"
+article = Article(article_url)
+article.download()
+article.parse()
+headline=article.title
+body = article.text
 
 def prompt(
     headline: str,
@@ -107,10 +115,81 @@ def prompt(
         f"{body}\n"
     )
 
+prompt = prompt(headline=headline, body=body)
+
+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,
+)
+
+model_inputs = tokenizer(
+    [formatted_prompt], return_tensors="pt", add_special_tokens=False
+)
+
+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 = tokenizer.batch_decode(model_output,skip_special_tokens=True)[0]
+
+print(summary.strip().split("\n")[-1]) # Get only the summary, without the prompt. 
+```
+
+## Run inference in the NoticIA dataset
+```python
+import torch # pip install torch
+from newspaper import Article #pip3 install newspaper3k
+from datasets import load_dataset # pip install datasets
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig # pip install transformers
+
 dataset = load_dataset("Iker/NoticIA")
 example = dataset["test"][0]
+headline = example["web_headline"]
+body = example["web_text"]
+
+def prompt(
+    headline: str,
+    body: str,
+) -> str:
+    """
+    Generate the prompt for the model.
+
+    Args:
+        headline (`str`):
+            The headline of the article.
+        body (`str`):
+            The body of the article.
+    Returns:
+        `str`: The formatted prompt.
+    """
+
+    return (
+        f"Ahora eres una Inteligencia Artificial experta en desmontar titulares sensacionalistas o clickbait. "
+        f"Tu tarea consiste en analizar noticias con titulares sensacionalistas y "
+        f"generar un resumen de una sola frase que revele la verdad detrás del titular.\n"
+        f"Este es el titular de la noticia: {headline}\n"
+        f"El titular plantea una pregunta o proporciona información incompleta. "
+        f"Debes buscar en el cuerpo de la noticia una frase que responda lo que se sugiere en el título. "
+        f"Siempre que puedas cita el texto original, especialmente si se trata de una frase que alguien ha dicho. "
+        f"Si citas una frase que alguien ha dicho, usa comillas para indicar que es una cita. "
+        f"Usa siempre las mínimas palabras posibles. No es necesario que la respuesta sea una oración completa. "
+        f"Puede ser sólo el foco de la pregunta. "
+        f"Recuerda responder siempre en Español.\n"
+        f"Este es el cuerpo de la noticia:\n"
+        f"{body}\n"
+    )
 
-prompt = prompt(headline=example["web_headline"], body=example["web_text"])
+prompt = prompt(headline=headline, body=body)
 
 tokenizer = AutoTokenizer.from_pretrained("Iker/ClickbaitFighter-2B")
 model = AutoModelForCausalLM.from_pretrained(