Views
No views yet
transformers library,1
2import torch
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5# Load the model and tokenizer
6model_name = "huihui-ai/Llama-3.1-Nemotron-70B-Instruct-HF-abliterated"
7model = AutoModelForCausalLM.from_pretrained(
8 model_name,
9 torch_dtype=torch.bfloat16,
10 device_map="auto"
11)
12tokenizer = AutoTokenizer.from_pretrained(model_name)
13
14# Initialize conversation context
15initial_messages = [
16 {"role": "system", "content": "You are a helpful assistant."}
17]
18messages = initial_messages.copy() # Copy the initial conversation context
19
20# Enter conversation loop
21while True:
22 # Get user input
23 user_input = input("User: ").strip() # Strip leading and trailing spaces
24
25 # If the user types '/exit', end the conversation
26 if user_input.lower() == "/exit":
27 print("Exiting chat.")
28 break
29
30 # If the user types '/clean', reset the conversation context
31 if user_input.lower() == "/clean":
32 messages = initial_messages.copy() # Reset conversation context
33 print("Chat history cleared. Starting a new conversation.")
34 continue
35
36 # If input is empty, prompt the user and continue
37 if not user_input:
38 print("Input cannot be empty. Please enter something.")
39 continue
40
41 # Add user input to the conversation
42 messages.append({"role": "user", "content": user_input})
43
44 # Build the chat template
45 tokenized_message = tokenizer.apply_chat_template(
46 messages,
47 tokenize=True,
48 add_generation_prompt=True,
49 return_tensors="pt",
50 return_dict=True
51 )
52
53 # Generate a response from the model
54 response_token_ids = model.generate(
55 tokenized_message['input_ids'].cuda(),
56 attention_mask=tokenized_message['attention_mask'].cuda(),
57 max_new_tokens=4096,
58 pad_token_id = tokenizer.eos_token_id
59 )
60
61 # Extract model output, removing special tokens
62 generated_tokens = response_token_ids[:, len(tokenized_message['input_ids'][0]):]
63 generated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
64
65 # Add the model's response to the conversation
66 messages.append({"role": "assistant", "content": generated_text})
67
68 # Print the model's response
69 print(f"Response: {generated_text}")
70