You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
import llamahf
|
|
import os
|
|
import torch
|
|
from transformers import StoppingCriteria, StoppingCriteriaList
|
|
|
|
# # to save memory use bfloat16
|
|
# torch.set_default_dtype(torch.bfloat16)
|
|
|
|
MODEL = 'decapoda-research/llama-7b-hf'
|
|
# MODEL = 'decapoda-research/llama-13b-hf'
|
|
# MODEL = 'decapoda-research/llama-30b-hf'
|
|
# MODEL = 'decapoda-research/llama-65b-hf'
|
|
|
|
if os.path.exists('./trained'):
|
|
MODEL = './trained'
|
|
|
|
tokenizer = llamahf.LLaMATokenizer.from_pretrained(MODEL)
|
|
model = llamahf.LLaMAForCausalLM.from_pretrained(MODEL, low_cpu_mem_usage=True)
|
|
model.to('cpu')
|
|
|
|
|
|
class StoppingCriteriaSub(StoppingCriteria):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, stops=[]):
|
|
if input_ids[0][-1] == 13:
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
ctx = """A dialog, where User interacts with AI. AI is helpful, kind, obedient, honest, and knows its own limits.
|
|
User: Hello, AI.
|
|
AI: Hello! How can I assist you today?
|
|
"""
|
|
|
|
while True:
|
|
print('-' * 40)
|
|
print(ctx.rstrip("\n"))
|
|
prompt = input(f'User: ')
|
|
if ctx != "":
|
|
ctx = ctx + "User: " + prompt + "\n"
|
|
else:
|
|
ctx = prompt + "\n"
|
|
|
|
ctx = (ctx[-1920:]) if len(ctx) >= 2048 else ctx
|
|
|
|
if len(ctx.strip()) > 0:
|
|
batch = tokenizer(ctx, return_tensors="pt")
|
|
result = model.generate(batch["input_ids"].cpu(),
|
|
do_sample=True,
|
|
top_k=50,
|
|
max_length=2048,
|
|
top_p=0.95,
|
|
temperature=1.0,
|
|
stopping_criteria=StoppingCriteriaList([StoppingCriteriaSub()]),
|
|
# repetition_penalty=1.17
|
|
)
|
|
decoded = tokenizer.decode(result[0])
|
|
ctx = decoded + "\n"
|