from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
import torch

model_name = "sudy-super/Contrail-200m-64k"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.bfloat16)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

if torch.cuda.is_available():
    model = model.to("cuda")

prompt = "AIによって私達の暮らしは、"

with torch.no_grad():
  token_ids = tokenizer.encode(prompt, return_tensors="pt")
  output_ids = model.generate(
      input_ids=token_ids.to(model.device),
      min_new_tokens=10,
      max_new_tokens=100,
      do_sample=True,
      temperature=0.7,
      streamer=streamer,
  )

Author

Rakuto Suda

Citations

@article{jiang2023mistral},
  title={Mistral 7B},
  author={Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, L{\'e}lio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timoth{\'e}e Lacroix, William El Sayed},
  journal={arXiv preprint arXiv:2310.06825},
  year={2023}
}