Views
No views yet

1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
3from vigogne.preprocess import generate_inference_chat_prompt
4
5model_name_or_path = "bofenghuang/vigogne-7b-chat"
6tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="right", use_fast=False)
7model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float16, device_map="auto")
8
9user_query = "Expliquez la différence entre DoS et phishing."
10prompt = generate_inference_chat_prompt([[user_query, ""]], tokenizer=tokenizer)
11input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)
12input_length = input_ids.shape[1]
13
14generated_outputs = model.generate(
15 input_ids=input_ids,
16 generation_config=GenerationConfig(
17 temperature=0.1,
18 do_sample=True,
19 repetition_penalty=1.0,
20 max_new_tokens=512,
21 ),
22 return_dict_in_generate=True,
23)
24generated_tokens = generated_outputs.sequences[0, input_length:]
25generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
26print(generated_text)