116 lines
3.8 KiB
Python
Executable File
116 lines
3.8 KiB
Python
Executable File
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
import re
|
|
|
|
class GPT2Node:
|
|
def __init__(self):
|
|
self.models = {}
|
|
self.tokenizer = None
|
|
self.rephrasers = ["microsoft/Promptist"]
|
|
pass
|
|
|
|
"""
|
|
A node for generating text using GPT-2, with the ability to specify the model via Hugging Face model identifier.
|
|
|
|
Attributes
|
|
----------
|
|
See base class.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
Defines input types for the GPT-2 node, including a model identifier.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"text": ("STRING", {"multiline": True}),
|
|
"model_repo": ("STRING", {
|
|
"default": "Gustavosta/MagicPrompt-Stable-Diffusion",
|
|
"multiline": False
|
|
}),
|
|
"temperature": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.1,
|
|
"max": 2.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "STRING",)
|
|
RETURN_NAMES = ("CONDITIONING","STRING",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "conditioning"
|
|
|
|
def generate(self, clip, text, model_repo, temperature):
|
|
"""
|
|
Generates text based on the provided prompt using a specified GPT-2 model.
|
|
|
|
Parameters
|
|
----------
|
|
text : str
|
|
The text prompt to generate text from.
|
|
model_repo : str
|
|
The Hugging Face model identifier for the GPT-2 model to use.
|
|
temperature : float
|
|
The temperature to use for generating text, controlling randomness.
|
|
|
|
Returns
|
|
-------
|
|
tuple
|
|
A tuple containing the generated text.
|
|
"""
|
|
if self.tokenizer == None:
|
|
self.tokenizer = AutoTokenizer.from_pretrained('gpt2')
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
self.tokenizer.padding_side = "left"
|
|
num_return_sequences = 1
|
|
max_length = 100 # max: 512
|
|
|
|
# Load the model based on the model_repo
|
|
if model_repo not in self.models:
|
|
self.models[model_repo] = AutoModelForCausalLM.from_pretrained(model_repo)
|
|
self.models[model_repo].eval() # for inference
|
|
|
|
text = text.strip()
|
|
if model_repo in self.rephrasers:
|
|
text = text+" Rephrase:"
|
|
input_ids = self.tokenizer(text, return_tensors='pt').input_ids
|
|
eos_id = self.tokenizer.eos_token_id
|
|
|
|
max_length = min(max_length + len(input_ids), 1024) # Adjusting max_length to account for prompt length
|
|
|
|
# Generate text
|
|
outputs = self.models[model_repo].generate(
|
|
input_ids,
|
|
max_length=max_length,
|
|
num_return_sequences=num_return_sequences,
|
|
temperature=temperature,
|
|
use_cache=True,
|
|
do_sample=False,
|
|
no_repeat_ngram_size=2,
|
|
eos_token_id=eos_id,
|
|
pad_token_id=eos_id,
|
|
length_penalty=-1.0,
|
|
)
|
|
|
|
output_texts = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
|
|
if model_repo in self.rephrasers:
|
|
new_prompt = re.sub(r'^.*? Rephrase:', '', output_texts[0]).strip()
|
|
else:
|
|
new_prompt = output_texts[0].strip()
|
|
|
|
tokens = clip.tokenize(new_prompt) # Tokenize the prompt
|
|
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) # Encode the tokens
|
|
return ([[cond, {"pooled_output": pooled}]], new_prompt,)
|
|
|
|
|
|
# Update the NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS to include the GPT2Node
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GPT2Node": GPT2Node
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GPT2Node": "GPT-2 Text Generator"
|
|
} |