KeraCare / drug_name_extraction_v2x0

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