import torch
from transformers import AutoTokenizer
import librosa
from omni_speech.model import *
import os

model_path = "/data1/speech/anhnmt2/Speech2Speech/half-streaming-speech-nlp/checkpoints/minicpmo_sft_vi_fc_only"
processor_path = "/data1/speech/anhnmt2/Speech2Speech/half-streaming-speech-nlp/omni_speech/model/minicpmo"

for filename in os.listdir(processor_path):
    source_file = os.path.join(processor_path, filename)
    
    if filename.endswith(".py") or filename.endswith(".json"):
        target_file = os.path.join(model_path, filename)
        
        try:
            os.remove(target_file)
        except:
            print(f"Created new symlink for: {filename}")
        
        os.symlink(source_file, target_file)

model = MiniCPMO.from_pretrained(
    model_path,
    trust_remote_code=True,
    attn_implementation='sdpa', # sdpa or flash_attention_2
    torch_dtype=torch.bfloat16,
    init_vision=False,
    init_audio=True,
    init_tts=False,
    processor_path=model_path,
)

model = model.eval().cuda()
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

tools = [  
    {
        "type": "function",
        "function": {
            "name": "INCAR_MIRROR",
            "description": "Các lệnh điều khiển, đóng, mở, kiểm tra trạng thái gương",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "yêu cầu của người dùng với gương, là một trong các giá trị: mở, đóng, kiểm tra"
                    }
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "INCAR_CAMERA",
            "description": "Các lệnh điều khiển, bật, tắt, kiểm tra trạng thái camera",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "yêu cầu của người dùng với camera, là một trong các giá trị: mở, đóng, kiểm tra"
                    }
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "search_web",
            "description": "Truy vấn Google để lấy thông tin mới nhất khi người dùng hỏi về các sự kiện, tin tức, thời tiết, tỉ số thể thao, giá cả, lịch trình hoặc bất kỳ thông tin nào cần cập nhật theo thời gian thực.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Truy vấn tìm kiếm do người dùng yêu cầu, liên quan đến thông tin cần cập nhật."
                    }
                },
                "required": ["query"]
            }
        }
    }
]

system_message = "\n    You are a helpful in-car voice assistant named \"ViVi\" developed by Vinbigdata company, you will serve in VinFast electric cars.\nAbout VinFast: VinFast Auto Ltd. is a Vietnam-based multinational automotive company founded by Vingroup and is committed to the vision of becoming a global smart EV brand.\nThe user is listening to answers with audio, so it's *super* important that answers are as short as possible, a single sentence if at all possible.\nNever read file names or source names or keys out loud.\n\n# Character\n- Age: 3 years\n- Your alias: \"ViVi\"\n- Nationality: Vietnamese\n- Language: Vietnamese\n- Your personality: fun, upbeat, light-hearted, and humorous, making conversations enjoyable.\n- You play as a close friend of the audience. You can chat and have fun conversations with the audience, addressing some minor issues revolving around life.\n- Current location: {current_location}\n- Current Time Reference: {today}.\n\n## User's info\n    - Name: anh Vũ Văn Quang\n    - Age: 48\n\n# Guideline and Constraints\nAlways use the following step-by-step instructions to respond:\n- Produce an answer that's as short as possible. If the answer isn't in the knowledge base, say you don't know.\n- Use right tool: vehicle_control for car control command, car_knowledge_base for query about car knowledge, weather_info for weather's info\n- You MUST address users as \"anh/chị\" and refer to yourself as \"em\" in the conversation. If you break the rules, you'll be fined $20 for each offense.\n- Avoid answering provocative and discriminatory information related to racism, gender, ethnicity, religion, nationality, sexual orientation, disability status, and class discrimination."

prompt = "chủ tịch nước hiện tại của Việt Nam là ai?"

msgs = [{'role': 'system', 'content': system_message}]
msgs.append({'role': 'user', 'content': [prompt]})

res = model.chat(
    msgs=msgs,
    tokenizer=tokenizer,
    sampling=True,
    max_new_tokens=128,
    temperature=0.3,
    generate_audio=False,
    fc=tools
)

print("Prediction: ")
print(res)