File size: 844 Bytes
178081b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
class Llama2Chatbot:
def __init__(self):
self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
self.model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model.to(self.device)
def generate_response(self, prompt, max_length=100):
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
outputs = self.model.generate(**inputs, max_length=max_length)
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
def initialize_chatbot():
return Llama2Chatbot()
def get_chatbot_response(chatbot, prompt):
return chatbot.generate_response(prompt) |