Spaces:
Runtime error
Runtime error
File size: 2,732 Bytes
7b715b6 | 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 | import os
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
def get_model_and_tokenizer(device, CHECKPOINTS_DIR, WEIGHTS_NAME, model_type):
config = AutoConfig.from_pretrained(model_type)
tokenizer = AutoTokenizer.from_pretrained(model_type, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
model_type,
from_tf=bool(".ckpt" in model_type),
config=config,
)
model.resize_token_embeddings(len(tokenizer))
model.gradient_checkpointing_enable()
model.to(device)
model_file = os.path.join(CHECKPOINTS_DIR, WEIGHTS_NAME)
try:
checkpoint = torch.load(model_file)
except:
checkpoint = torch.load(model_file, map_location=device)
model.load_state_dict(checkpoint)
return tokenizer, model
def generate_prompt(title=""):
return f"<BOS> TOPIC: {title} COMMENT: "
def extract_comment_from_prompt(text):
if type(text) == str:
starts = text.rfind("COMMENT: ") + len("COMMENT: ")
ends = text.find(" <EOS>")
result = text[starts:ends]
else:
text = pd.Series(text)
starts = text.str.find("COMMENT: ") + len("COMMENT: ")
ends = text.str.find(" <EOS>")
result = [sentence[start:end] for (sentence, start, end) in zip(text, starts, ends)]
return result
def generate_comment(prompt, tokenizer, model, entry_length=30, top_p=0.8, temperature=1.):
filter_value = -float("Inf")
model.eval()
with torch.no_grad():
generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
for i in range(entry_length):
outputs = model(generated, labels=generated)
loss, logits = outputs[:2]
logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0
indices_to_remove = sorted_indices[sorted_indices_to_remove]
logits[:, indices_to_remove] = filter_value
next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
generated = torch.cat((generated, next_token), dim=1)
if next_token in tokenizer.encode("<EOS>"):
break
output_list = list(generated.squeeze().numpy())
output_text = f"{tokenizer.decode(output_list)} <EOS>"
return output_text |