pytorch

working in progress

from transformers import AutoTokenizer, AutoModelForCausalLM
import transformers
import torch

model = "Terry0320/Kestrel"

tokenizer = AutoTokenizer.from_pretrained(model)
model = AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True)
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto",
)
sequences = pipeline(
   "Generelativa is obsessed with general relativity and believes that it is the greatest theory in the world. He thinks that any other theory, including special relativity, is simply not worth mentioning.\nHoward: Hello, Generelativa!\nGenerelativa:",
    max_length=200,
    do_sample=True,
    top_k=10,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
)
for seq in sequences:
    print(f"Result: {seq['generated_text']}")