from transformers import AutoProcessor, AutoModelForImageTextToText, AutoConfig
from PIL import Image
import torch
import torch.nn.functional as F
from peft import PeftModel
device = "cuda:0"
model_id = "google/gemma-3-12b-it"
peft_model_path = "historyHulk/ModiTrans-12B-Gemma-Teacher"
model = AutoModelForImageTextToText.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map=device
)
model = PeftModel.from_pretrained(
model,
peft_model_path,
device_map=device,
torch_dtype=torch.bfloat16
)
image = Image.open("<Modi Script Image Preprocessed as in Dataset>").convert("RGB").resize((1024,512))
processor = AutoProcessor.from_pretrained(model_id)
messages = [
{
"role": "user",
"content": [
{
"type": "image",
"image":image
},
{
"type": "text",
"text": "Translitrate the following Modi script to Devnagri script."
},
],
},
{
"role": "assistant",
"content": [
{
"type": "text",
},
],
},
]
inputs = processor.apply_chat_template(
messages, add_generation_prompt=True, tokenize=True,
return_dict=True, return_tensors="pt"
).to(model.device, dtype=torch.bfloat16)
input_len = inputs["input_ids"].shape[-1]
pixel_values = inputs['pixel_values']
pixel_values = pixel_values.to(dtype=model.dtype, device=model.device)
model.eval()
with torch.no_grad():
input_ids = inputs["input_ids"]
attention_masks = inputs["attention_mask"]
pixel_values=pixel_values
while True:
outputs = model(
input_ids=input_ids,
attention_mask=attention_masks,
pixel_values=pixel_values,
)
logits = outputs.logits[:,-1,:]
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
input_ids = torch.cat([input_ids, next_token], dim=-1)
attention_masks = torch.cat([attention_masks, torch.ones_like(next_token)], dim=-1)
if next_token.item() == processor.tokenizer.eos_token_id or input_ids.shape[1] >= 350:
break
generation = input_ids[:,input_len:][0]
generated_text = processor.decode(generation, skip_special_tokens=True)
print("\n\n\n")
print(generated_text)