# 所需包
flash_attn==2.5.8
numpy==1.24.4
Pillow==10.3.0
Requests==2.31.0
torch==2.3.0
torchvision==0.18.0
transformers==4.40.2

from PIL import Image
import requests
from transformers import AutoModelForCausalLM, AutoProcessor

class Phi3VisionModel:
    def __init__(self, model_id="microsoft/Phi-3-vision-128k-instruct", device="cuda"):
        """
        Initialize the Phi3VisionModel with the specified model ID and device.
        
        Args:
            model_id (str): The identifier of the pre-trained model from Hugging Face's model hub.
            device (str): The device to load the model on ("cuda" for GPU or "cpu").
        """
        self.model_id = model_id
        self.device = device
        self.model = self.load_model()  # 在初始化期間加載模型
        self.processor = self.load_processor()  # 在初始化期間加載處理器
    
    def load_model(self):
        """
        Load the pre-trained language model with causal language modeling capabilities.
        
        Returns:
            model (AutoModelForCausalLM): The loaded model.
        """
        print("Loading model...")
        # 使用自動(dòng)設(shè)備映射和數(shù)據(jù)類型調(diào)整加載模型
        return AutoModelForCausalLM.from_pretrained(
            self.model_id, 
            device_map="auto",  # 自動(dòng)將模型映射到適當(dāng)?shù)脑O(shè)備
            torch_dtype="auto",  # 根據(jù)設(shè)備使用適當(dāng)?shù)?torch 數(shù)據(jù)類型
            trust_remote_code=True,  # 允許執(zhí)行自定義代碼以加載模型
            _attn_implementation='flash_attention_2'  # 使用優(yōu)化的注意力實(shí)現(xiàn)
        ).to(self.device)  # 將模型移動(dòng)到指定的設(shè)備
    
    def load_processor(self):
        """
        Load the processor associated with the model for processing inputs and outputs.
        
        Returns:
            processor (AutoProcessor): The loaded processor for handling text and images.
        """
        print("Loading processor...")
        # 使用 trust_remote_code=True 加載處理器，以處理任何自定義處理邏輯
        return AutoProcessor.from_pretrained(self.model_id, trust_remote_code=True)
    
    def predict(self, image_url, prompt):
        """
        Perform a prediction using the model given an image and a prompt.
        
        Args:
            image_url (str): The URL of the image to be processed.
            prompt (str): The textual prompt that guides the model's generation.
        
        Returns:
            response (str): The generated response from the model.
        """
        # 從提供的 URL 加載圖像
        image = Image.open(requests.get(image_url, stream=True).raw)
        
        # 為模型格式化輸入提示模板
        prompt_template = f"<|user|>\n<|image_1|>\n{prompt}<|end|>\n<|assistant|>\n"
        
        # 處理輸入，將提示和圖像轉(zhuǎn)換為張量格式
        inputs = self.processor(prompt_template, [image], return_tensors="pt").to(self.device)
        
        # 設(shè)置模型響應(yīng)生成的生成參數(shù)
        generation_args = {
            "max_new_tokens": 500,  # 生成的最大令牌數(shù)
            "temperature": 0.7,     # 生成的采樣溫度，以增加多樣性
            "do_sample": False      # 禁用采樣以獲得確定性輸出
        }
        print("Generating response...")
        # 使用模型生成輸出 ID，跳過輸入令牌
        output_ids = self.model.generate(**inputs, **generation_args)
        output_ids = output_ids[:, inputs['input_ids'].shape[1]:]  # Ignore the input prompt in the output
        
        # 解碼生成的輸出令牌以獲得響應(yīng)文本
        response = self.processor.batch_decode(output_ids, skip_special_tokens=True)[0]
        return response

# 初始化模型
phi_model = Phi3VisionModel()

# 示例預(yù)測
image_url = "https://example.com/sample_image.png"  # 示例圖像的 URL  
prompt = "Extract the data in json format."  # 指導(dǎo)模型的提示  
response = phi_model.predict(image_url, prompt)  # 獲取模型的響應(yīng) 

print("Response:", response)  # 打印生成的響應(yīng)

{
  "Type/Type": "P",
  "Country code/Code du pays": "UTO",
  "Passport Number/N° de passeport": "L898902C3",
  "Surname/Nom": "ERIKSSON",
  "Given names/Prénoms": "ANNA MARIA",
  "Nationality/Nationalité": "UTOPIAN",
  "Date of Birth/Date de naissance": "12 AUGUST/AOUT 74",
  "Personal No./N° personnel": "Z E 184226 B",
  "Sex/Sexe": "F",
  "Place of birth/Lieu de naissance": "ZENITH",
  "Date of issue/Date de délivrance": "16 APR/AVR 07",
  "Authority/Autorité": "PASSPORT OFFICE",
  "Date of expiry/Date d'expiration": "15 APR/AVR 12",
  "Holder's signature/Signature du titulaire": "anna maria eriksson",
  "Passport/Passeport": "P<UTOERIKSSON<<ANNA<MARIA<<<<<<<<<<<<<<<<<<<<<<<L898902C36UT07408122F1204159ZE184226B<<<<10"
}

{
  "passport": {
    "issuingCountry": "Netherlands",
    "issuingAuthority": "Koninkrijk der Nederlanden",
    "passportNumber": "SPEC12014",
    "issuingDate": "09 MAR 2014",
    "expiryDate": "09 MAR 2024",
    "holder": {
      "gender": "F",
      "nationality": "Netherlands",
      "placeOfBirth": "SPECIMEN",
      "sex": "WF",
      "firstNames": [
        "Willem",
        "Lieselotte"
      ]
    },
    "physicalDescription": {
      "height": "1.75 m",
      "hairColor": "gray",
      "hairLength": "short"
    },
    "issuingOffice": "Burg. van Stad en Dorp",
    "issuingDateAsInt": "14032014",
    "expiryDateAsInt": "14032024",
    "fieldsExtracted": [
      {
        "code": "NL",
        "dateOfBirth": "10 MAR 1965",
        "dateOfIssue": "09 MAR 2014",
        "dateOfExpiry": "09 MAR 2024",
        "firstNames": [
          "Willem",
          "Lieselotte"
        ],
        "nationality": "Netherlands",
        "passportNumber": "SPEC12014",
        "placeOfBirth": "SPECIMEN",
        "sex": "WF"
      }
    ]
  }
}