|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
import torch |
|
|
|
model_name = "EQUES/TinyDeepSeek-1.5B" |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto") |
|
|
|
|
|
input_text = "Explain large language models in simple terms." |
|
input_ids = tokenizer(input_text, return_tensors="pt").input_ids |
|
|
|
|
|
output = model.generate(input_ids, max_length=100) |
|
generated_text = tokenizer.decode(output[0], skip_special_tokens=True) |
|
|
|
print(generated_text) |
|
|