Spaces:
Runtime error
Runtime error
Commit History
-
src/chatbots/gptjbot.py
from transformers import GPTJForCausalLM, AutoTokenizer import os import torch device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") class DialoGPT: def __init__( self, model_name: str = "EleutherAI/gpt-j-6B", local_path="./models/gpt-j-6B", ): if not os.path.exists(local_path): GPTJForCausalLM.from_pretrained(model_name).save_pretrained( local_path, revision="float16", torch_dtype=torch.float16, ) AutoTokenizer.from_pretrained(model_name).save_pretrained(local_path) self.model = GPTJForCausalLM.from_pretrained( local_path, revision="float16", torch_dtype=torch.float16, ) self.tokenizer = AutoTokenizer.from_pretrained(local_path) def __call__(self, inputs: str) -> str: input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids.to(device) generated_ids = se…13a8ecd -
Create src/chatbots/blenderbot.py
c5529b3