Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_id = "CATIE-AQ/LMF2-350M_french_summary"
4model = AutoModelForCausalLM.from_pretrained(
5 model_id,
6 device_map="auto",
7 torch_dtype="bfloat16",
8 trust_remote_code=True,
9# attn_implementation="flash_attention_2" <- uncomment on compatible GPU
10)
11tokenizer = AutoTokenizer.from_pretrained(model_id)
12
13prompt = """Résume l'article suivant :\n""" + "you_text_to_summarize"
14tokenizer.padding_side = "left"
15tokenizer.truncation_side = "right"
16
17# Apply the chat template to prepare the input
18input_ids = tokenizer.apply_chat_template(
19 [{"role": "user", "content": prompt}],
20 add_generation_prompt=True,
21 return_tensors="pt",
22 tokenize=True,
23).to(model.device)
24
25# Generate the output from the model
26output = model.generate(
27 input_ids,
28 do_sample=True,
29 temperature=0.3,
30 min_p=0.15,
31 repetition_penalty=1.05,
32 max_new_tokens=2048,
33)
34
35summary_text = tokenizer.decode(
36 output[0][input_ids.shape[1]:], # Slice the output tensor
37 skip_special_tokens=True # Skip special tokens for a cleaner output
38)
39
40print(summary_text)