Views
No views yet
1import torch
2import transformers
3from peft import PeftModel
4from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
5import json
6
7peft_model_id = "raicrits/Llama_NewsTagger_XRECO"
8base_model_id = "meta-llama/Llama-3.1-8B-Instruct"
9
10
11base_model = AutoModelForCausalLM.from_pretrained(base_model_id,
12 quantization_config=quantization_config,
13 device_map='auto')
14model = PeftModel.from_pretrained(base_model, peft_model_id)
15model.config.use_cache = True
16
17
18tokenizer = AutoTokenizer.from_pretrained(peft_model_id, add_eos_token=False)
19tokenizer.pad_token = tokenizer.eos_token
20terminators = [
21 tokenizer.eos_token_id,
22 tokenizer.convert_tokens_to_ids("<|eot_id|>")
23]
24
25def generate_prompt_tag(title, subtitle, text):
26 messages = [
27 {"role": "system", "content": "You are a multilanguage AI assistant aimed to suggest tags in news articles."},
28 {"role": "user", "content": f"""Analyze the following news article and assign a list of representative tags to the content, returning always at least 3.
29
30Title: {title}
31
32Subtitle: {subtitle}
33
34Text: {text}
35
36The response has to be given in the following format "[<tag1>, <tag2>,.. ]", where the <tag>s are the tags identified and written in the same language of the article. Do no add any further text."""}
37]
38 prompt = tokenizer.apply_chat_template(
39 messages,
40 add_generation_prompt=True,
41 return_tensors="pt",
42 return_dict=True).to(model.device)
43 return prompt
44
45prompt = generate_prompt_tag(<title>, <subtitle>, <text>) #INSERT NEWS ARTICLE TITLE, SUBTITLE AND TEXT HERE
46
47outputs = model.generate(
48 input_ids=prompt['input_ids'], attention_mask=prompt['attention_mask'],
49 eos_token_id=terminators,
50 pad_token_id=tokenizer.eos_token_id
51 )
52response = outputs[0][prompt['input_ids'].shape[-1]:]
53
54tags = tokenizer.decode(response, skip_special_tokens=True).replace("\n","").replace("[","").replace("]", "").replace("'","").replace('"','')
55