Model Card for KeraCare/drug_name_extraction_v2x0
This model is a fine-tuned version of the GLM-OCR model, trained for drug name extraction from prescription images. It was fine-tuned on a custom dataset of prescription images and corresponding drug name labels.
Usage
To use this model for inference, you can load it using the Hugging Face Transformers library:
from transformers import AutoProcessor, AutoModelForImageTextToText
import time
import torch
processor = AutoProcessor.from_pretrained("KeraCare/drug_name_extraction_v2x0")
model = AutoModelForImageTextToText.from_pretrained("KeraCare/drug_name_extraction_v2x0")
image_path = "sample-images/test_image.jpg"
messages = [
{
"role": "user",
"content": [
{
"type": "image",
"image": image_path
},
{
"type": "text",
"text": "Extract drug names from the image in json format with the following format: {\"drug_names\": [\"drug_name1\", \"drug_name2\", ...]}"
}
],
}
]
inputs = processor.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
enable_thinking=False,
)
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
model.to(device)
model.eval()
inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
# Per-sample input lengths (handles padding correctly for batch_size > 1)
input_lengths = inputs["attention_mask"].sum(dim=1).tolist()
generated_ids = model.generate(
**inputs,
max_new_tokens=128
)
output_ids = generated_ids[0][int(input_lengths[0]):]
generated_text = processor.decode(output_ids, skip_special_tokens=True)
print("Generated Text:")
print(generated_text)
Training Details
- Base Model: GLM-OCR
- Fine-tuning Dataset: Custom dataset of prescription images and drug name labels