kadirnar
e21711a049
- Updated model selection in image-to-image, mask generation, object detection, video classification, zero-shot image classification, zero-shot object detection, document question answering, feature extraction, image-to-text, image-text-to-text, visual question answering, conversational, fill mask, question answering, sentence similarity, table question answering, zero-shot classification, text classification, text generation, and token classification pipelines to use a default string instead of a model list. - Removed unused keypoint detection, summarization, text2text generation, and translation pipelines. - Added utility nodes for type conversion: IntToString, StringToInt, FloatToString, and StringToFloat. - Updated all pipeline calls to include trust_remote_code=True for enhanced security. - Improved requirements by adding timm and opencv-python for better functionality.
24 lines
962 B
Python
24 lines
962 B
Python
from transformers import pipeline as hf_pipeline
|
|
|
|
class TextGenerationPipeline:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"model_name": ("STRING", {"default": "gpt2"}),
|
|
"max_new_tokens": ("INT", {"default": 50, "min": 1, "max": 2048}),
|
|
"temperature": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 2.0, "step": 0.05}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("generated_text",)
|
|
FUNCTION = "run_generation"
|
|
CATEGORY = "Transformers/NLP/TextGeneration"
|
|
|
|
def run_generation(self, prompt, model_name, max_new_tokens, temperature):
|
|
pipe = hf_pipeline("text-generation", model=model_name, trust_remote_code=True)
|
|
result = pipe(prompt, max_new_tokens=max_new_tokens, temperature=temperature)
|
|
return (result[0]["generated_text"],)
|