# 部署Llama 3.2-3B文本模型
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载模型
model_name = "meta-llama/Llama-3.2-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 对话示例
def chat_with_llama(message, history=[]):
    messages = history + [{"role": "user", "content": message}]

    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            input_ids,
            max_new_tokens=512,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )

    response = tokenizer.decode(
        outputs[0][input_ids.shape[-1]:],
        skip_special_tokens=True
    )

    return response

# 使用示例
response = chat_with_llama("请解释什么是边缘计算？")
print(response)

# 部署Llama 3.2-11B-Vision多模态模型
from transformers import MllamaForConditionalGeneration, AutoProcessor
from PIL import Image
import torch

# 加载视觉模型
model_name = "meta-llama/Llama-3.2-11B-Vision-Instruct"
processor = AutoProcessor.from_pretrained(model_name)
model = MllamaForConditionalGeneration.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 图像分析函数
def analyze_image(image_path, question):
    # 加载图像
    image = Image.open(image_path)

    # 准备输入
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": question}
            ]
        }
    ]

    # 处理输入
    input_text = processor.apply_chat_template(
        messages,
        add_generation_prompt=True
    )
    inputs = processor(
        image,
        input_text,
        return_tensors="pt"
    ).to(model.device)

    # 生成回答
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=1000,
            do_sample=True,
            temperature=0.7
        )

    response = processor.decode(
        output[0][inputs['input_ids'].shape[-1]:],
        skip_special_tokens=True
    )

    return response

# 使用示例
response = analyze_image(
    "document.jpg",
    "请提取这个文档中的关键信息"
)
print(response)

# 使用ONNX Runtime进行移动端部署
import onnxruntime as ort
import numpy as np

class MobileLlama:
    def __init__(self, model_path):
        # 加载ONNX模型
        self.session = ort.InferenceSession(
            model_path,
            providers=['CPUExecutionProvider']
        )

    def generate(self, input_ids, max_length=512):
        # 移动端推理逻辑
        outputs = self.session.run(
            None,
            {'input_ids': input_ids.astype(np.int64)}
        )
        return outputs[0]

# 部署到移动设备
mobile_model = MobileLlama("llama-3.2-1b-mobile.onnx")