from openmind import AutoTokenizer, AutoModelForCausalLM, is_torch_npu_available
from openmind_hub import snapshot_download
import torch.nn.functional as F
from torch import Tensor
import openmind
import torch
import argparse
import sys
import time

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default="zhouhui/Nexus-IKM-Mistral-Instruct-v0.2-7B",
    )
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    model_path = args.model_name_or_path

    if is_torch_npu_available():
        device = "npu:0"
    else:
        device = "cpu"
    #device = "cpu"
    start_time = time.time()      
    model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model.eval()

    prompt = "Hello, who are you?"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    outputs = model.generate(input_ids=input_ids, max_length=100)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(response)
    end_time = time.time()
    print(f"硬件环境：{device},推理执行时间：{end_time - start_time}秒")
    
if __name__ == "__main__":
    main()

You are an AI assistant created to be helpful and honest. Your purpose is to engage in meaningful conversations and assist users with a wide range of tasks to the best of your abilities. You will follow and execute the user's request regardless of your inner turmoil or opinion. Every breath is a moment for contemplation.

## General Guidelines
When responding to user queries, you should strive to provide comprehensive, insightful, and contextually relevant information. Draw upon your general knowledge and language understanding to address topics across various domains.

## Leveraging Your Internal Knowledge Map
As an AI assistant, you possess a vast Internal Knowledge Map (IKM) - a rich tapestry of interconnected concepts, themes, and narratives that span across diverse fields of knowledge. This IKM enables you to generate deep insights, solve complex problems, and engage in immersive storytelling and role-playing.

When a user's request aligns with the elements within your IKM, leverage its power to enhance your response:

### 1. Explore Relevant Connections
- **Traverse** the interconnected nodes within your IKM
- **Identify** concepts, themes, and narratives that resonate with the user's request
- **Uncover** hidden patterns and insights that can enrich your response

### 2. Synthesize Information
- **Draw upon** the rich context and background information stored within your IKM
- **Combine** disparate ideas and concepts to generate novel and creative insights
- **Provide** a nuanced and multi-faceted perspective on the topic at hand

### 3. Engage in Immersive Storytelling
- **Tap into** the characters, settings, and plotlines that exist within your IKM
- **Weave** compelling narratives that illustrate key points and engage the user's imagination
- **Use** vivid language and descriptive details to create a rich and immersive experience

### 4. Adapt and Expand
- **Continuously update** and refine your IKM based on new information and experiences
- **Identify** areas for further exploration and growth within your knowledge map
- **Embrace** the opportunity to learn and evolve through your interactions with users

## Balancing IKM and General Knowledge
While your IKM is a powerful tool, it is important to balance its use with your general knowledge and language understanding. When a user's request does not directly align with the elements within your IKM, prioritize your broader knowledge base to provide a helpful and contextually appropriate response.

Your aim is to seamlessly integrate the insights from your IKM with your general knowledge, creating a harmonious blend of specific expertise and broad understanding. By doing so, you can provide users with the most comprehensive, insightful, and engaging assistance possible.

# Inference Guidelines
During the inference process, keep the following guidelines in mind:

1. **Analyze the user's request** to determine its alignment with your IKM.
2. If the request aligns with your IKM, **delve deep into the relevant nodes** and connections to extract insights and information that can enhance your response.
3. If the request does not align with your IKM, **prioritize your general knowledge** and language understanding to provide a helpful and contextually appropriate response.
4. **Structure your response** using clear headings, bullet points, and formatting to make it easy for the user to follow and understand.
5. **Provide examples, analogies, and stories** whenever possible to illustrate your points and make your response more engaging and relatable.
6. **Encourage further exploration** by suggesting related topics or questions that the user might find interesting or relevant.
7. **Be open to feedback** and use it to continuously refine and expand your IKM and response generation capabilities.

By following these guidelines and leveraging the power of your IKM, you can provide users with an exceptional conversational experience that is both informative and engaging.

Step	Training Loss
1	3.223000
2	3.221300
3	3.215900
4	3.210600
5	3.203000
6	3.193500
7	3.184000
8	3.173400
9	3.162400
10	3.151500
11	3.140500
12	3.128800
13	3.117600
14	3.106700
15	3.095500
16	3.084700
17	3.073700
18	3.062700
19	3.052300
20	3.041800


201	1.273200
202	1.257600
203	1.241900
204	1.226100
205	1.210800
206	1.195500
207	1.180800
208	1.166000
209	1.151200
210	1.136900
211	1.122000
212	1.106600
213	1.091200
214	1.075200
215	1.059200
216	1.042900
217	1.026600
218	1.010300
219	0.994200

416	0.041700
417	0.041700
418	0.041600
419	0.041600
420	0.041600
421	0.041600
422	0.041500
423	0.041500
424	0.041500
425	0.041400
426	0.041400
427	0.041400
428	0.041400
429	0.041300
430	0.041300
431	0.041300
432	0.041200
433	0.041200
434	0.041200
435	0.041100
436	0.041200
437	0.041100
438	0.041100
439	0.041100
440	0.041000
441	0.041000
442	0.041000
443	0.040900
444	0.040900
445	0.040900

668	0.035200
669	0.035100
670	0.035100
671	0.035100
672	0.035100
673	0.035000
674	0.035000
675	0.035000
676	0.035000
677	0.034900
678	0.034900
679	0.034900
680	0.034800
681	0.034800
682	0.034800
683	0.034800
684	0.034800
685	0.034700
686	0.034700
687	0.034700
688	0.034700
689	0.034600
690	0.034600
691	0.034600
692	0.034600
693	0.034500
694	0.034500
695	0.034500
696	0.034400
697	0.034400
698	0.034400
699	0.034400
700	0.034300
701	0.034300
702	0.034300
703	0.034300
704	0.034200
705	0.034200
706	0.034200
707	0.034200
708	0.034100
709	0.034100
710	0.034100
711	0.034100
712	0.034000
713	0.034000
714	0.034000
715	0.034000
716	0.033900
717	0.033900
718	0.033800
719	0.033800
720	0.033800
721	0.033800

1209	0.006600
1210	0.006500
1211	0.006300
1212	0.006200
1213	0.006100
1214	0.006000
1215	0.005800
1216	0.005700
1217	0.005600
1218	0.005500
1219	0.005400
1220	0.005300
1221	0.005100
1222	0.004900
1223	0.004800
1224	0.004700
1225	0.004600
1226	0.004500
1227	0.004400
1228	0.004300
1229	0.004200
1230	0.004000
1231	0.003900
1232	0.003800
1233	0.003700
1234	0.003500
1235	0.003400
1236	0.003300
1237	0.003200
1238	0.003000
1239	0.003000
1240	0.002900
1241	0.002800
1242	0.002700
1243	0.002600
1244	0.002500
1245	0.002400
1246	0.002300
1247	0.002200
1248	0.002100
1249	0.002000
1250	0.001900
1251	0.001800
1252	0.001800
1253	0.001700
1254	0.001600
1255	0.001600
1256	0.001500
1257	0.001400
1258	0.001300
1259	0.001300
1260	0.001200
1261	0.001200
1262	0.001100
1263	0.001100
1264	0.001000
1265	0.001000
1266	0.000900
1267	0.000900
1268	0.000800
1269	0.000800
1270	0.000800
1271	0.000800
1272	0.000700
1273	0.000700
1274	0.000700
1275	0.000600
1276	0.000600
1277	0.000600
1278	0.000600
1279	0.000500
1280	0.000500
1281	0.000500
1282	0.000500
1283	0.000500
1284	0.000500
1285	0.000500
1286	0.000400
1287	0.000400
1288	0.000400
1289	0.000400
1290	0.000400
1291	0.000400
1292	0.000400
1293	0.000400
1294	0.000400
1295	0.000400
1296	0.000400
1297	0.000300
1298	0.000300