File size: 1,460 Bytes
1580220
 
 
 
 
 
 
 
 
 
 
 
 
 
b03c35f
a71c2cf
 
1580220
a71c2cf
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load model and tokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
    "universitytehran/PersianMind-v1.0",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    device_map={"": device},
)
tokenizer = AutoTokenizer.from_pretrained("universitytehran/PersianMind-v1.0")

# Conversation template
TEMPLATE = "{context}\nYou: {prompt}\nPersianMind: "
CONTEXT = "به انگلیسی ترجمه کن"
prompt = "به انگلیسی ترجمه کن: وزیر امور خارجه ایران در نشست خبری اعلام کرد که مذاکرات هسته‌ای به دلیل اختلافات اساسی درباره اجرای تعهدات متوقف شده است. او تأکید کرد که ایران همچنان به دنبال یک توافق عادلانه و متوازن است، اما طرف‌های مقابل باید از سیاست‌های دوگانه خود دست بردارند"

model_input = TEMPLATE.format(context=CONTEXT, prompt=prompt)
input_tokens = tokenizer(model_input, return_tensors="pt").to(device)
generate_ids = model.generate(**input_tokens, max_new_tokens=512, do_sample=False, repetition_penalty=1.1)
model_output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
response = model_output[len(model_input):]
print(response)