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