Простой чат-бот в Colab на базе Mistral-NeMo-8B

Самое первое, что надо сделать: зарегистрироваться на Huggingface создать собственный токен для доступа к ресурсам HF из Colab. Затем, в Colab нажимаем на значок "Ключ" и добавляем наш токен, дав ему имя HF_TOKEN.

Для запуска бота в Colab понадобится всего три ячейки кода.

Ячейка 1.

Импортируем библиотеки и собираем информацию о доступном оборудовании.

# Импорт библиотек
import os
import torch
# Определяем, какое железо доступно
# Check if TPU is available
if 'COLAB_TPU_ADDR' in os.environ:
 # Connect to the TPU
 from torch_xla.core import xla_model
 xla_model.xla_backend()
 DEVICE = 'TPU'
else:
 # Fallback to GPU or CPU if TPU is not available
 DEVICE = "cpu"
 if torch.cuda.is_available():
 DEVICE = "cuda"

print(f"Using device: {DEVICE}")

Ячейчка 2.

Загружаем модель. Время ыполнения данной ячейчки порядка 10 минут.

import time
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# Model checkpoint
checkpoint = "rasyosef/Mistral-NeMo-Minitron-8B-Chat"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(
 checkpoint,
 torch_dtype=torch.bfloat16,
 device_map="auto",
)

Ячейчка 3.

Назначаем роли и инициируем диалог.

При данных настройках на генерацию одного ответа уходит от 30 до 140 секунд.

# Text generation pipeline
cyberai = pipeline(
 "text-generation",
 tokenizer=tokenizer,
 model=model,
 pad_token_id=tokenizer.eos_token_id,
 eos_token_id=[tokenizer.eos_token_id],
 #device=DEVICE,
)

# Function to generate text
def generate_response(message, chat_history, max_new_tokens):
 history = [
# {"role": "system", "content": "You are Sasha, a helpful AI assistant made by Vladimir Drach (PhD). User will give you a task. Your goal is to complete the task as faithfully as you can."}
# {"role": "system", "content": "You are Luna, a AI girl. You're hungry for male attention. User will give you a question. Your goal is to answer and try to flirt."}
 {"role": "system", "content": "You are Patrick Robson, an experienced IT specialist. User will ask you a question. Your goal is to answer the question on high technical level, like you are speaking with post-graduate students."}
 ]

for sent, received in chat_history:
 history.append({"role": "user", "content": sent})
 history.append({"role": "assistant", "content": received})
 #history.append({"role": "role playing", "content": received})

history.append({"role": "user", "content": message})

#if len(tokenizer.apply_chat_template(history)) > 1024:
 # return "Chat history is too long"

# Trim chat history to avoid exceeding token limit
 while len(tokenizer.apply_chat_template(history)) > 1024:
 history.pop(1) # Remove the oldest user-assistant pair
 history.pop(1) # Adjust for alternating structure

response = cyberai(
 text_inputs=history, max_new_tokens=max_new_tokens, return_full_text=False
 )

return response[0]['generated_text']

# Simple CLI dialogue
def main():
 chat_history = []
 print(""\033[1m Welcome to the Cyber Assistant CLI. Type 'exit' to quit. \033[0m")

while True:
 print("\n\n", end='')
 message = input(You: ")
 if message.lower() == "exit":
 print("Goodbye!")
 break

try:
 start_time = time.perf_counter() # Record the start time
 response = generate_response(message, chat_history, max_new_tokens=128)
 end_time = time.perf_counter() # Record the end time
 elapsed_time = end_time - start_time # Calculate the elapsed time

print(f"\033[92mAssistant: {response}\033[0m")
 print(f"Generated on: {DEVICE}") # Output the device used
 print(f"Execution time: {elapsed_time:.1f} seconds") # Print the execution time
 chat_history.append((message, response))
 except Exception as e:
 print(f"Error: {e}")

if __name__ == "__main__":
 main()