Views
No views yet
viena_data/examples/pretrain_offline.jsonlconfigs/viena_tiny.yaml1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4model_id = "vietrix/viena-tiny-demo-pretrain"
5tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
6model = AutoModelForCausalLM.from_pretrained(
7 model_id,
8 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
9 device_map="auto",
10)
11
12prompt = "Viena la gi?
13"
14inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
15output = model.generate(**inputs, max_new_tokens=128, do_sample=True, temperature=0.7, top_p=0.9)
16print(tokenizer.decode(output[0], skip_special_tokens=True))