能不能把它封装成一个类

当然可以！将ChitGPT封装成一个类可以更方便地使用和管理。下面是一个简单的示例代码：

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

class ChitGPT:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
        self.model = GPT2LMHeadModel.from_pretrained("gpt2")
        self.model.to(self.device)
        self.model.eval()

    def generate_response(self, input_text):
        input_ids = self.tokenizer.encode(input_text, return_tensors="pt").to(self.device)
        with torch.no_grad():
            output = self.model.generate(input_ids, max_length=100, num_return_sequences=1)
        response = self.tokenizer.decode(output[0], skip_special_tokens=True)
        return response

使用示例：

chit_gpt = ChitGPT()
input_text = "你好"
response = chit_gpt.generate_response(input_text)
print(response)

这是一个简单的封装示例，你可以根据自己的需求进行扩展和优化。