Llama-3.2-3B-Reasoning-Vi-Medical-LoRA

import os
from huggingface_hub import login

# Set the Hugging Face API token
os.environ["HUGGINGFACEHUB_API_TOKEN"] = "<your_huggingface_token>"

# # Initialize API
login(os.environ.get("HUGGINGFACEHUB_API_TOKEN"))

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

device = "cuda" if torch.cuda.is_available() else "cpu"

# Define model and LoRA adapter paths
base_model_name = "unsloth/Llama-3.2-3B-Instruct"
lora_adapter_name = "danhtran2mind/Llama-3.2-3B-Reasoning-Vi-Medical-LoRA"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model_name)

# Load base model with optimized settings
model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.float16,  # Use FP16 for efficiency
    device_map=device,
    trust_remote_code=True
)

# Apply LoRA adapter
model = PeftModel.from_pretrained(model, lora_adapter_name)

# Set model to evaluation mode
model.eval()

inference_prompt_style = """Bên dưới là một hướng dẫn mô tả một tác vụ, đi kèm với một thông tin đầu vào để cung cấp thêm ngữ cảnh.
Hãy viết một phản hồi để hoàn thành yêu cầu một cách phù hợp.
Trước khi trả lời, hãy suy nghĩ cẩn thận về câu hỏi và tạo một chuỗi suy nghĩ từng bước để đảm bảo phản hồi logic và chính xác.

### Instruction:
Bạn là một chuyên gia y tế có kiến thức chuyên sâu về lập luận lâm sàng, chẩn đoán và lập kế hoạch điều trị.
Vui lòng trả lời câu hỏi y tế sau đây.

### Question:
{}

### Response:
<think>
"""

# Define the question
question = ("Khi nghi ngờ bị loét dạ dày tá tràng nên đến khoa nào "
            "tại bệnh viện để thăm khám?")

seed = 42
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

inputs = tokenizer(
    [inference_prompt_style.format(question) + tokenizer.eos_token],
    return_tensors="pt"
).to(device)

outputs = model.generate(
    **inputs,
    max_new_tokens=2048,
    temperature=0.7,
    top_p=0.95,
    top_k=64,
)

response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print(response[0].split("### Response:")[1])