Views
No views yet
1!pip install peft
2!pip install transformers
3!pip install bitsandbytes
4!pip install accelerate1# You need a huggingface token that can access llama2
2from huggingface_hub import notebook_login
3notebook_login()1import torch
2from peft import PeftModel, PeftConfig
3from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
4device = "cuda" if torch.cuda.is_available() else "cpu"
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
13peft_model_id = "danjie/Chadgpt-Llama2-70b-conversation"
14config = PeftConfig.from_pretrained(peft_model_id)
15model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, device_map='cuda', quantization_config=bnb_config)
16tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
17
18# Load the Lora model
19model = PeftModel.from_pretrained(model, peft_model_id)1# Run this cell to start a new conversation
2conversation_history = []
3
4def format_conversation(conversation: list[str]) -> str:
5 formatted_conversation = ""
6
7 # Check if the conversation has more than two turns
8 if len(conversation) > 2:
9 # Process all but the last two turns
10 for i in range(len(conversation) - 2):
11 if i % 2 == 0:
12 formatted_conversation += "<Past User>" + conversation[i] + "\n"
13 else:
14 formatted_conversation += "<Past Assistant>" + conversation[i] + "\n"
15
16 # Process the last two turns
17 if len(conversation) >= 2:
18 formatted_conversation += "<User>" + conversation[-2] + "\n"
19 formatted_conversation += "<Assistant>" + conversation[-1]
20
21 return formatted_conversation
22
23def talk_with_llm(chat: str) -> str:
24 # Encode and move tensor into cuda if applicable.
25 conversation_history.append(chat)
26 conversation_history.append("")
27 conversation = format_conversation(conversation_history)
28
29 encoded_input = tokenizer(conversation, return_tensors='pt')
30 encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
31
32 output = model.generate(**encoded_input, max_new_tokens=256)
33 response = tokenizer.decode(output[0], skip_special_tokens=True)
34 response = response[len(conversation):]
35
36 conversation_history.pop()
37 conversation_history.append(response)
38 return response