Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4# ၁။ Device ကို အရင်သတ်မှတ်မယ် (GPU ရှိရင် cuda၊ မရှိရင် cpu)
5device = "cuda" if torch.cuda.is_available() else "cpu"
6
7model_name = "URajinda/Qwen2.5-MM-1.5B-Base"
8
9# ၂။ Model ကို load လုပ်ပြီး GPU ဆီ ပို့မယ် (.to(device) ထည့်လိုက်တာပါ)
10model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
11tokenizer = AutoTokenizer.from_pretrained(model_name)
12
13# ၃။ စာရိုက်မယ်
14prompt = "မြန်မာနိုင်ငံတွင်"
15
16# ၄။ Input တွေကိုလည်း GPU ဆီ ပို့မယ် (.to(device) ထည့်လိုက်တာပါ)
17inputs = tokenizer(prompt, return_tensors="pt").to(device)
18
19# ၅။ အဖြေထုတ်မယ်
20outputs = model.generate(
21 **inputs,
22 max_new_tokens=100,
23 temperature=0.7,
24 do_sample=True
25)
26
27print(tokenizer.decode(outputs[0], skip_special_tokens=True))