import tensorflow as tf
from transformers import AutoTokenizer, TFOPTForCausalLM
# Load the model and tokenizer
model = TFOPTForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
# Run Inference
prompt = "Are you conscious? Can you talk?"
# Tokenize the input text
input_ids = tokenizer(prompt, return_tensors='tf').input_ids
@tf.function
def generate():
return model.generate(input_ids, max_length=20)
# Run inference
outputs = generate()
# Decode the outputs
decoded_outputs = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
for result in decoded_outputs:
print(result)
Sample Output
I can talk, but I can't really think