Spaces:
Runtime error
Runtime error
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model_dir = "Den4ikAI/ebany_researcher" | |
| tokenizer = AutoTokenizer.from_pretrained(model_dir) | |
| tokenizer.add_special_tokens({'bos_token': '<s>', 'eos_token': '</s>', 'pad_token': '<pad>'}) | |
| model = AutoModelForCausalLM.from_pretrained(model_dir) | |
| model.to(device) | |
| model.eval() | |
| def generate(prompt): | |
| encoded_prompt = tokenizer.encode(prompt, return_tensors="pt").to(device) | |
| out = model.generate(encoded_prompt, max_length=50, do_sample=True, top_k=50, top_p=0.95, temperature=1.0, | |
| num_return_sequences=1) | |
| for i, tokens in enumerate(out.cpu().tolist(), start=1): | |
| tokens = tokens[encoded_prompt.shape[1]:] | |
| text = tokenizer.decode(tokens) | |
| return text | |