Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
2import torch
3
4# Load model and tokenizer
5tokenizer = AutoTokenizer.from_pretrained("AquilaX-AI/QnA")
6model = AutoModelForCausalLM.from_pretrained("AquilaX-AI/QnA")
7
8# Define the system prompt
9prompt = """
10<|im_start|>system\nYou are a helpful AI assistant named Securitron<|im_end|>
11"""
12
13# Initialize conversation history
14conversation_history = []
15
16# Set up device
17device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
18model.to(device)
19
20while True:
21 user_prompt = input("\nUser Question: ")
22 if user_prompt.lower() == 'break':
23 break
24
25 # Format the user's input
26 user = f"""<|im_start|>user
27{user_prompt}<|im_end|>
28<|im_start|>assistant"""
29
30 # Add the user's question to the conversation history
31 conversation_history.append(user)
32
33 # Keep only the last 2 exchanges (4 turns)
34 conversation_history = conversation_history[-5:]
35
36 # Build the full prompt
37 current_prompt = prompt + "\n".join(conversation_history)
38
39 # Tokenize the prompt
40 encodeds = tokenizer(current_prompt, return_tensors="pt", truncation=True).input_ids.to(device)
41
42 # Initialize TextStreamer for real-time token generation
43 text_streamer = TextStreamer(tokenizer, skip_prompt=True)
44
45 # Generate response with TextStreamer
46 response = model.generate(
47 input_ids=encodeds,
48 streamer=text_streamer,
49 max_new_tokens=512,
50 use_cache=True,
51 pad_token_id=151645,
52 eos_token_id=151645,
53 num_return_sequences=1
54 )
55
56 # Finalize conversation history with the assistant's response
57 conversation_history.append(tokenizer.decode(response[0]).split('<|im_start|>assistant')[-1].split('<|im_end|>')[0].strip() + "<|im_end|>")