Lsavints/swo_llm

1from threading import Thread2from typing import Iterator3 4import torch5from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer6 7from huggingface_hub import HfFolder8 9def save_huggingface_token(token: str):10    folder = HfFolder()11    folder.save_token(token)12save_huggingface_token("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")13 14model_id = 'Lsavints/swo_llm'15 16if torch.cuda.is_available():17    model = AutoModelForCausalLM.from_pretrained(18        model_id,19        torch_dtype=torch.float16,20        device_map='auto',21        use_auth_token=True22    )23else:24    model = None25tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True,26                                          device_map="auto"27                                          )28tokenizer.pad_token = tokenizer.eos_token29tokenizer.padding_side = "right"30 31 32def get_prompt(message: str, chat_history: list[tuple[str, str]],33               system_prompt: str) -> str:34    texts = [f'<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n']35    # The first user input is _not_ stripped36    do_strip = False37    for user_input, response in chat_history:38        user_input = user_input.strip() if do_strip else user_input39        do_strip = True40        texts.append(f'{user_input} [/INST] {response.strip()} </s><s>[INST] ')41    message = message.strip() if do_strip else message42    texts.append(f'{message} [/INST]')43    return ''.join(texts)44 45 46def get_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> int:47    prompt = get_prompt(message, chat_history, system_prompt)48    input_ids = tokenizer([prompt], return_tensors='np', add_special_tokens=False)['input_ids']49    return input_ids.shape[-1]50 51 52def run(message: str,53        chat_history: list[tuple[str, str]],54        system_prompt: str,55        max_new_tokens: int = 1024,56        temperature: float = 0.8,57        top_p: float = 0.95,58        top_k: int = 50) -> Iterator[str]:59    prompt = get_prompt(message, chat_history, system_prompt)60    inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')61 62    streamer = TextIteratorStreamer(tokenizer,63                                    timeout=10.,64                                    skip_prompt=True,65                                    skip_special_tokens=True)66    generate_kwargs = dict(67        inputs,68        streamer=streamer,69        max_new_tokens=max_new_tokens,70        do_sample=True,71        top_p=top_p,72        top_k=top_k,73        temperature=temperature,74        num_beams=1,75    )76    t = Thread(target=model.generate, kwargs=generate_kwargs)77    t.start()78 79    outputs = []80    for text in streamer:81        outputs.append(text)82        yield ''.join(outputs)83