Mr-Vicky-01/qwen-conversational-finetuned

INFERENCE

# Load model directly
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("Mr-Vicky-01/qwen-conversational-finetuned")
model = AutoModelForCausalLM.from_pretrained("Mr-Vicky-01/qwen-conversational-finetuned")

prompt = """
<|im_start|>system\nYou are a helpful AI assistant named Securitron<|im_end|>
"""

# Keep a list for the last one conversation exchanges
conversation_history = []

while True:
    user_prompt = input("\nUser Question: ")
    if user_prompt.lower() == 'break':
        break

    # Format the user's input
    user = f"""<|im_start|>user
{user_prompt}<|im_end|>
<|im_start|>assistant"""

    # Add the user's question to the conversation history
    conversation_history.append(user)

    # Ensure conversation starts with a user's input and keep only the last 2 exchanges (4 turns)
    conversation_history = conversation_history[-5:]

    # Build the full prompt
    current_prompt = prompt + "\n".join(conversation_history)

    # Tokenize the prompt
    encodeds = tokenizer(current_prompt, return_tensors="pt", truncation=True).input_ids

    # Move model and inputs to the appropriate device
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    inputs = encodeds.to(device)

    # Create an empty list to store generated tokens
    generated_ids = inputs

    # Start generating tokens one by one
    assistant_response = ""
    for _ in range(512):  # Specify a max token limit for streaming
        next_token = model.generate(
            generated_ids,
            max_new_tokens=1,
            pad_token_id=151644,
            eos_token_id=151645,
            num_return_sequences=1,
            do_sample=False,
            # top_k=5,
            # temperature=0.2,
            # top_p=0.90
        )
        
        generated_ids = torch.cat([generated_ids, next_token[:, -1:]], dim=1)
        token_id = next_token[0, -1].item()
        token = tokenizer.decode([token_id], skip_special_tokens=True)
        
        assistant_response += token
        print(token, end="", flush=True)

        if token_id == 151645:  # EOS token
            break

    conversation_history.append(f"{assistant_response.strip()}<|im_end|>")