Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
3
4model_id = "SalmanFaroz/Llama-2-7b-samsum"
5
6bnb_config = BitsAndBytesConfig(
7 load_in_4bit=True,
8 bnb_4bit_use_double_quant=True,
9 bnb_4bit_quant_type="nf4",
10 bnb_4bit_compute_dtype=torch.bfloat16
11)
12
13model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map="auto")
14
15tokenizer = AutoTokenizer.from_pretrained(model_id)
16tokenizer.pad_token = tokenizer.eos_token
17tokenizer.padding_side = "right"
18
19# Define the input prompt
20prompt = """
21Summarize the following conversation.
22
23### Input:
24Itachi: Kakashi, you must understand the gravity of the situation. The Akatsuki's plans are far more sinister than you can imagine.
25Kakashi: Itachi, I need more than vague warnings. Tell me what you know.
26Itachi: Very well. The Akatsuki seeks to capture Naruto for the power of the Nine-Tails sealed within him, but there's an even darker secret lurking within their goals.
27Kakashi: Darker than that? What are they truly after?
28Itachi: They're hunting the Tailed Beasts for a cataclysmic plan to reshape the world, and only we can stop them, together.
29
30### Summary:
31"""
32
33inputs = tokenizer(prompt, return_tensors='pt')
34output = tokenizer.decode(
35 model.generate(
36 inputs["input_ids"],
37 max_new_tokens=100,
38 )[0],
39 skip_special_tokens=True
40)
41
42print("Output:",output)
43