File size: 3,416 Bytes
3fe575d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
MODEL_PATH = "VDrontV1"
TEMPERATURE = 0.40
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
def load_model_and_tokenizer(model_path):
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=False)
model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
device_map="auto",
trust_remote_code=False
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# Получаем ID токена <|endoftext|> если он есть в словаре
endoftext_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
if endoftext_id != tokenizer.unk_token_id: # если токен существует
model.config.eos_token_id = endoftext_id
tokenizer.eos_token = "<|endoftext|>"
tokenizer.eos_token_id = endoftext_id
return model, tokenizer
def generate_stream(model, tokenizer, prompt, temperature=0.4, max_new_tokens=128):
# Text continuation generator without chat history
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
# Получаем eos_token_id для генерации
eos_token_id = tokenizer.eos_token_id
generation_kwargs = dict(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
do_sample=True,
top_p=0.95,
repetition_penalty=1.1,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=eos_token_id,
streamer=streamer,
)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
for new_text in streamer:
# Проверяем, не содержит ли текст <|endoftext|>
if "<|endoftext|>" in new_text:
# Обрезаем до маркера
new_text = new_text.split("<|endoftext|>")[0]
yield new_text
break
yield new_text
thread.join()
def interactive_chat(model, tokenizer, temperature):
print(f"Simple text continuation (temp={temperature})")
print("Enter the beginning of the text, and the model will continue it.")
print("Commands: 'exit' or 'quit' — exit.")
while True:
try:
user_input = input("\nYou: ").strip()
except (KeyboardInterrupt, EOFError):
print("\nGoodbye!")
break
if user_input.lower() in ["exit", "quit"]:
print("Goodbye!")
break
if not user_input:
continue
print() # empty line before continuation
for token in generate_stream(model, tokenizer, "<|user|> " + user_input + " <|assistant|>",
temperature=temperature):
print(token, end="", flush=True)
print() # finish the line
if __name__ == "__main__":
model, tokenizer = load_model_and_tokenizer(MODEL_PATH)
interactive_chat(model, tokenizer, temperature=TEMPERATURE) |