import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaListfrom threading import Thread checkpoint_path = "Vladniag/Bormolom" # укажите свой репозиторий model = AutoModelForCausalLM.from_pretrained( checkpoint_path, device_map="auto", torch_dtype=torch.float16)tokenizer = AutoTokenizer.from_pretrained(checkpoint_path, use_fast=True)if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token class StopOnTokens(StoppingCriteria): def __init__(self, stop_token_ids_list): self.stop_token_ids_list = stop_token_ids_list def __call__(self, input_ids, scores, **kwargs): for stop_ids in self.stop_token_ids_list: if len(input_ids[0]) >= len(stop_ids): if input_ids[0][-len(stop_ids):].tolist() == stop_ids: return True return False stop_words = ["</result>"]stop_ids_list = [tokenizer.encode(w, add_special_tokens=False) for w in stop_words]stopping_criteria = StoppingCriteriaList([StopOnTokens(stop_ids_list)]) def generate_response(prompt: str, max_new_tokens: int = 512) -> str: processed = "▁" + prompt.lstrip().replace(" ", "▁").replace("\n", "<newline>") inputs = tokenizer(processed, return_tensors="pt", add_special_tokens=False).to(model.device) streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False) kwargs = dict( inputs, streamer=streamer, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.3, top_p=0.5, repetition_penalty=1.3, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, stopping_criteria=stopping_criteria, ) thread = Thread(target=model.generate, kwargs=kwargs) thread.start() full = "" for chunk in streamer: decoded = chunk.replace("▁", " ").replace("<newline>", "\n").replace("\\n", "\n") if "</result>" in decoded: full += decoded.split("</result>")[0] break full += decoded return full