Every Inference gives the prompt as part of output also, any way to remove that?
#26
by
vermanic
- opened
Hey, i have this general problem of any model on HF outputting with the input prompt always, any way to exclude that?
as I just need the output.
Code:
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
checkpoint = "HuggingFaceH4/starchat-beta"
device = "cuda" if torch.cuda.is_available() else "cpu" # "cuda:X" for GPU usage or "cpu" for CPU usage
class Model:
def __init__(self):
print("Running in " + device)
self.tokenizer = AutoTokenizer.from_pretrained(checkpoint)
self.model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map='auto')
def infer(self, input_text, token_count):
inputs = self.tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = self.model.generate(inputs, max_new_tokens=token_count)
return self.tokenizer.decode(outputs[0])
Also, max_new_tokens means the amount of tokens with which I want the model to respond with, right?
vermanic
changed discussion title from
Every Inference gives the prompt as part of output also, any way to fix this?
to Every Inference gives the prompt as part of output also, any way to remove that?
Resolved by:
return self.tokenizer.decode(outputs[0])[len(input_text):]
vermanic
changed discussion status to
closed