|
import os
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
|
|
|
|
|
import torch
|
|
import numpy as np
|
|
|
|
from transformers import AutoModelForSequenceClassification, AutoModelForMultipleChoice, AutoModel
|
|
from transformers import TrainingArguments, Trainer
|
|
from transformers import AutoTokenizer
|
|
from transformers import DataCollatorWithPadding
|
|
from datasets import load_dataset
|
|
from datasets import load_metric
|
|
import torchsnooper
|
|
|
|
model_name = "THUDM/chatglm-6b"
|
|
|
|
if __name__ == "__main__":
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
|
model = AutoModel.from_pretrained(model_name, trust_remote_code=True).half().cuda()
|
|
response, history = model.chat(tokenizer, "你好", history=[])
|
|
print(response)
|
|
response, history = model.chat(tokenizer, "晚上睡不着应该怎么办", history=history)
|
|
print(response)
|
|
|
|
|
|
|
