# 部署Phi-3 Mini模型
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载模型
model_name = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

# 对话函数
def chat_with_phi3(message, system_prompt="You are a helpful AI assistant."):
    messages = [
        {"role": "system", "content": system_prompt},
        {"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=1000,
            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_phi3("请解释量子计算的基本原理")
print(response)

# 部署Phi-3 Mini 128K长上下文版本
model_name = "microsoft/Phi-3-mini-128k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

# 长文档处理函数
def process_long_document(document, question):
    messages = [
        {
            "role": "system",
            "content": "你是一个专业的文档分析助手，能够处理长文档并回答相关问题。"
        },
        {
            "role": "user",
            "content": f"文档内容：\n{document}\n\n问题：{question}"
        }
    ]

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

    # 检查输入长度
    if input_ids.shape[1] > 128000:
        print(f"警告：输入长度 {input_ids.shape[1]} 超过128K限制")
        return "文档过长，请分段处理"

    with torch.no_grad():
        outputs = model.generate(
            input_ids,
            max_new_tokens=2000,
            do_sample=True,
            temperature=0.3,
            top_p=0.9
        )

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

    return response

# 使用示例
long_doc = """这里是一个很长的文档内容..."""
question = "请总结文档的主要观点"
response = process_long_document(long_doc, question)
print(response)

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

class MobilePhi3:
    def __init__(self, model_path):
        # 配置ONNX Runtime
        self.session = ort.InferenceSession(
            model_path,
            providers=[
                'CPUExecutionProvider',
                # 'CoreMLExecutionProvider',  # iOS
                # 'NNAPIExecutionProvider',   # Android
            ]
        )

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

# 量化优化
from transformers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# 加载量化模型
model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3-mini-4k-instruct",
    quantization_config=quantization_config,
    device_map="auto",
    trust_remote_code=True
)