Spaces:
Runtime error
Runtime error
| 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 |