internlm/internlm3-8b-instruct

1"""This script refers to the dialogue example of streamlit, the interactive2generation code of chatglm2 and transformers.3 4We mainly modified part of the code logic to adapt to the5generation of our model.6Please refer to these links below for more information:7    1. streamlit chat example:8        https://docs.streamlit.io/knowledge-base/tutorials/build-conversational-apps9    2. chatglm2:10        https://github.com/THUDM/ChatGLM2-6B11    3. transformers:12        https://github.com/huggingface/transformers13Please run with the command `streamlit run path/to/web_demo.py14    --server.address=0.0.0.0 --server.port 7860`.15Using `python path/to/web_demo.py` may cause unknown problems.16"""17 18# isort: skip_file19import copy20import re21import warnings22from dataclasses import asdict, dataclass23from typing import Callable, List, Optional24 25import streamlit as st26import torch27from torch import nn28 29from transformers.generation.utils import LogitsProcessorList30from transformers.utils import logging31 32from transformers import AutoTokenizer, AutoModelForCausalLM  # isort: skip33 34logger = logging.get_logger(__name__)35st.set_page_config(layout='wide')36 37 38@dataclass39class GenerationConfig:40    # this config is used for chat to provide more diversity41    max_length: int = 3276842    top_p: float = 0.843    temperature: float = 0.844    do_sample: bool = True45    repetition_penalty: float = 1.00546 47 48@torch.inference_mode()49def generate_interactive(50    model,51    tokenizer,52    prompt,53    generation_config: Optional[GenerationConfig] = None,54    logits_processor: Optional[LogitsProcessorList] = None,55    prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,56    additional_eos_token_id: Optional[int] = None,57    **kwargs,58):59    inputs = tokenizer([prompt], padding=True, return_tensors='pt')60    input_length = len(inputs['input_ids'][0])61    for k, v in inputs.items():62        inputs[k] = v.to(model.device)63    input_ids = inputs['input_ids']64    _, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]65    if generation_config is None:66        generation_config = model.generation_config67    generation_config = copy.deepcopy(generation_config)68    generation_config._eos_token_tensor = generation_config.eos_token_id69    model_kwargs = generation_config.update(**kwargs)70    if generation_config.temperature == 0.0:71        generation_config.do_sample = False72    eos_token_id = generation_config.eos_token_id73    if isinstance(eos_token_id, int):74        eos_token_id = [eos_token_id]75    if additional_eos_token_id is not None:76        eos_token_id.append(additional_eos_token_id)77    has_default_max_length = kwargs.get('max_length') is None and generation_config.max_length is not None78    if has_default_max_length and generation_config.max_new_tokens is None:79        warnings.warn(80            f"Using 'max_length''s default \81                ({repr(generation_config.max_length)}) \82                to control the generation length. "83            'This behaviour is deprecated and will be removed from the \84                config in v5 of Transformers -- we'85            ' recommend using `max_new_tokens` to control the maximum \86                length of the generation.',87            UserWarning,88        )89    elif generation_config.max_new_tokens is not None:90        generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length91        if not has_default_max_length:92            logger.warn(  # pylint: disable=W490293                f"Both 'max_new_tokens' (={generation_config.max_new_tokens}) "94                f"and 'max_length'(={generation_config.max_length}) seem to "95                "have been set. 'max_new_tokens' will take precedence. "96                'Please refer to the documentation for more information. '97                '(https://huggingface.co/docs/transformers/main/'98                'en/main_classes/text_generation)',99                UserWarning,100            )101 102    if input_ids_seq_length >= generation_config.max_length:103        input_ids_string = 'input_ids'104        logger.warning(105            f'Input length of {input_ids_string} is {input_ids_seq_length}, '106            f"but 'max_length' is set to {generation_config.max_length}. "107            'This can lead to unexpected behavior. You should consider'108            " increasing 'max_new_tokens'."109        )110 111    # 2. Set generation parameters if not already defined112    logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()113    logits_processor = model._get_logits_processor(114        generation_config=generation_config,115        input_ids_seq_length=input_ids_seq_length,116        encoder_input_ids=input_ids,117        prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,118        logits_processor=logits_processor,119    )120    unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)121    while True:122        model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)123        # forward pass to get next token124        outputs = model(125            **model_inputs,126            return_dict=True,127            output_attentions=False,128            output_hidden_states=False,129        )130 131        next_token_logits = outputs.logits[:, -1, :]132 133        # pre-process distribution134        next_token_scores = logits_processor(input_ids, next_token_logits)135 136        # sample137        probs = nn.functional.softmax(next_token_scores, dim=-1)138        if generation_config.do_sample:139            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)140        else:141            next_tokens = torch.argmax(probs, dim=-1)142 143        # update generated ids, model inputs, and length for next step144        input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)145        unfinished_sequences = unfinished_sequences.mul((min(next_tokens != i for i in eos_token_id)).long())146 147        output_token_ids = input_ids[0].cpu().tolist()148        output_token_ids = output_token_ids[input_length:]149        for each_eos_token_id in eos_token_id:150            if output_token_ids[-1] == each_eos_token_id:151                output_token_ids = output_token_ids[:-1]152        response = tokenizer.decode(output_token_ids)153 154        yield response155        # stop when each sentence is finished156        # or if we exceed the maximum length157        if unfinished_sequences.max() == 0:158            break159 160 161def on_btn_click():162    del st.session_state.messages163    del st.session_state.deepthink_messages164 165 166def postprocess(text, add_prefix=True, deepthink=False):167    text = re.sub(r'\\\(|\\\)', r'$', text)168    text = re.sub(r'\\\[|\\\]', r'$$', text)169    if add_prefix:170        text = (':red[[Deep Thinking]]\n\n' if deepthink else ':blue[[Normal Response]]\n\n') + text171    return text172 173 174@st.cache_resource175def load_model():176    model_path = 'internlm/internlm3-8b-instruct'177    model = AutoModelForCausalLM.from_pretrained(178        model_path, trust_remote_code=True, device_map='auto', torch_dtype=torch.bfloat16179    )180    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)181    return model, tokenizer182 183 184def prepare_generation_config():185    with st.sidebar:186        max_length = st.slider('Max Length', min_value=8, max_value=32768, value=32768)187        top_p = st.slider('Top P', 0.0, 1.0, 0.8, step=0.01)188        temperature = st.slider('Temperature', 0.0, 1.0, 0.7, step=0.01)189        radio = st.radio('Inference Mode', ['Normal Response', 'Deep Thinking'], key='mode')190        st.button('Clear Chat History', on_click=on_btn_click)191 192    st.session_state['inference_mode'] = radio193    generation_config = GenerationConfig(max_length=max_length, top_p=top_p, temperature=temperature)194 195    return generation_config196 197 198user_prompt = '<|im_start|>user\n{user}<|im_end|>\n'199robot_prompt = '<|im_start|>assistant\n{robot}<|im_end|>\n'200cur_query_prompt = '<|im_start|>user\n{user}<|im_end|>\n\201    <|im_start|>assistant\n'202 203 204def combine_history(prompt, deepthink=False, start=0, stop=None):205    if stop is None:206        stop = len(st.session_state.messages)207    elif stop < 0:208        stop = len(st.session_state.messages) + stop209    messages = []210    for idx in range(start, stop):211        message, deepthink_message = st.session_state.messages[idx], st.session_state.deepthink_messages[idx]212        if deepthink:213            if deepthink_message['content'] is not None:214                messages.append(deepthink_message)215            else:216                messages.append(message)217        else:218            if message['content'] is not None:219                messages.append(message)220            else:221                messages.append(deepthink_message)222    meta_instruction = (223        'You are InternLM (书生·浦语), a helpful, honest, '224        'and harmless AI assistant developed by Shanghai '225        'AI Laboratory (上海人工智能实验室).'226    )227    if deepthink:228        meta_instruction += """You are an expert mathematician with extensive experience in mathematical competitions. You approach problems through systematic thinking and rigorous reasoning. When solving problems, follow these thought processes:229## Deep Understanding230Take time to fully comprehend the problem before attempting a solution. Consider:231- What is the real question being asked?232- What are the given conditions and what do they tell us?233- Are there any special restrictions or assumptions?234- Which information is crucial and which is supplementary?235## Multi-angle Analysis236Before solving, conduct thorough analysis:237- What mathematical concepts and properties are involved?238- Can you recall similar classic problems or solution methods?239- Would diagrams or tables help visualize the problem?240- Are there special cases that need separate consideration?241## Systematic Thinking242Plan your solution path:243- Propose multiple possible approaches244- Analyze the feasibility and merits of each method245- Choose the most appropriate method and explain why246- Break complex problems into smaller, manageable steps247## Rigorous Proof248During the solution process:249- Provide solid justification for each step250- Include detailed proofs for key conclusions251- Pay attention to logical connections252- Be vigilant about potential oversights253## Repeated Verification254After completing your solution:255- Verify your results satisfy all conditions256- Check for overlooked special cases257- Consider if the solution can be optimized or simplified258- Review your reasoning process259Remember:2601. Take time to think thoroughly rather than rushing to an answer2612. Rigorously prove each key conclusion2623. Keep an open mind and try different approaches2634. Summarize valuable problem-solving methods2645. Maintain healthy skepticism and verify multiple times265Your response should reflect deep mathematical understanding and precise logical thinking, making your solution path and reasoning clear to others.266When you're ready, present your complete solution with:267- Clear problem understanding268- Detailed solution process269- Key insights270- Thorough verification271Focus on clear, logical progression of ideas and thorough explanation of your mathematical reasoning. Provide answers in the same language as the user asking the question, repeat the final answer using a '\\boxed{}' without any units, you have [[8192]] tokens to complete the answer.272"""  # noqa: E501273    total_prompt = f'<s><|im_start|>system\n{meta_instruction}<|im_end|>\n'274    for message in messages:275        cur_content = message['content']276        if message['role'] == 'user':277            cur_prompt = user_prompt.format(user=cur_content)278        elif message['role'] == 'robot':279            cur_prompt = robot_prompt.format(robot=cur_content)280        else:281            raise RuntimeError282        total_prompt += cur_prompt283    total_prompt = total_prompt + cur_query_prompt.format(user=prompt)284    return total_prompt285 286 287def main():288    # torch.cuda.empty_cache()289    print('load model begin.')290    model, tokenizer = load_model()291    print('load model end.')292 293    user_avator = 'assets/user.png'294    robot_avator = 'assets/robot.png'295 296    st.title('InternLM3-8B-Instruct')297 298    generation_config = prepare_generation_config()299 300    def render_message(msg, msg_idx, deepthink):301        if msg['content'] is None:302            real_prompt = combine_history(303                st.session_state.messages[msg_idx - 1]['content'], deepthink=deepthink, stop=msg_idx - 1304            )305            placeholder = st.empty()306            for cur_response in generate_interactive(307                model=model,308                tokenizer=tokenizer,309                prompt=real_prompt,310                additional_eos_token_id=92542,311                **asdict(generation_config),312            ):313                placeholder.markdown(postprocess(cur_response, deepthink=deepthink) + '▌')314            placeholder.markdown(postprocess(cur_response, deepthink=deepthink))315            msg['content'] = cur_response316            torch.cuda.empty_cache()317        else:318            st.markdown(postprocess(msg['content'], deepthink=deepthink))319 320    # Initialize chat history321    if 'messages' not in st.session_state:322        st.session_state.messages = []323    if 'deepthink_messages' not in st.session_state:324        st.session_state.deepthink_messages = []325 326    # Display chat messages from history on app rerun327    for idx, (message, deepthink_message) in enumerate(328        zip(st.session_state.messages, st.session_state.deepthink_messages)329    ):330        with st.chat_message(message['role'], avatar=message.get('avatar')):331            if message['role'] == 'user':332                st.markdown(postprocess(message['content'], add_prefix=False))333            else:334                if st.toggle('compare', key=f'compare_{idx}'):335                    cols = st.columns(2)336                    if st.session_state['inference_mode'] == 'Deep Thinking':337                        with cols[1]:338                            render_message(deepthink_message, idx, True)339                        with cols[0]:340                            render_message(message, idx, False)341                    else:342                        with cols[0]:343                            render_message(message, idx, False)344                        with cols[1]:345                            render_message(deepthink_message, idx, True)346                else:347                    if st.session_state['inference_mode'] == 'Deep Thinking':348                        if deepthink_message['content'] is not None:349                            st.markdown(postprocess(deepthink_message['content'], deepthink=True))350                        else:351                            st.markdown(postprocess(message['content']))352                    else:353                        if message['content'] is not None:354                            st.markdown(postprocess(message['content']))355                        else:356                            st.markdown(postprocess(deepthink_message['content'], deepthink=True))357 358    # Accept user input359    if prompt := st.chat_input('What is up?'):360        # Display user message in chat message container361        with st.chat_message('user', avatar=user_avator):362            st.markdown(postprocess(prompt, add_prefix=False))363        real_prompt = combine_history(prompt, deepthink=st.session_state['inference_mode'] == 'Deep Thinking')364        # Add user message to chat history365        st.session_state.messages.append({'role': 'user', 'content': prompt, 'avatar': user_avator})366        st.session_state.deepthink_messages.append({'role': 'user', 'content': prompt, 'avatar': user_avator})367 368        with st.chat_message('robot', avatar=robot_avator):369            st.toggle('compare', key=f'compare_{len(st.session_state.messages)}')370            message_placeholder = st.empty()371            for cur_response in generate_interactive(372                model=model,373                tokenizer=tokenizer,374                prompt=real_prompt,375                additional_eos_token_id=92542,376                **asdict(generation_config),377            ):378                # Display robot response in chat message container379                message_placeholder.markdown(380                    postprocess(cur_response, deepthink=st.session_state['inference_mode'] == 'Deep Thinking') + '▌'381                )382            message_placeholder.markdown(383                postprocess(cur_response, deepthink=st.session_state['inference_mode'] == 'Deep Thinking')384            )385        # Add robot response to chat history386        response, deepthink_response = (387            (None, cur_response) if st.session_state['inference_mode'] == 'Deep Thinking' else (cur_response, None)388        )389        st.session_state.messages.append(390            {391                'role': 'robot',392                'content': response,  # pylint: disable=undefined-loop-variable393                'avatar': robot_avator,394            }395        )396        st.session_state.deepthink_messages.append(397            {398                'role': 'robot',399                'content': deepthink_response,400                'avatar': robot_avator,401            }402        )403        torch.cuda.empty_cache()404 405 406main()407