import time
start_time = time.time()

from optimum.intel import OVModelForCausalLM
import openvino.properties as props
import openvino.properties.hint as hints
import openvino.properties.streams as streams
from transformers import AutoTokenizer, GenerationConfig
import torch

# Configure optimization settings
ov_config = {
    hints.performance_mode(): hints.PerformanceMode.LATENCY,
    streams.num(): "1",
    props.cache_dir(): "",
    "DYNAMIC_QUANTIZATION_GROUP_SIZE": "32",  # For CPU optimization
    "KV_CACHE_PRECISION": "u8"  # For memory optimization
}

# Load model with optimizations
model_name_or_id = "AI4Chem/CHEMLLM-2b-1_5"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_id, trust_remote_code=True)

# Convert and optimize model
model = OVModelForCausalLM.from_pretrained(
    model_name_or_id,
    export=True,
    ov_config=ov_config,
    device="CPU",
    # load_in_8bit=True,  # Int8 quantization for OpenVINO is applied differently
    torch_dtype=torch.float16,
    trust_remote_code=True  # Include if the model uses custom code
)

end_time = time.time()
execution_time = end_time - start_time
print(f"Model load and optimization time: {execution_time} seconds")

/opt/app-root/lib64/python3.9/site-packages/transformers/cache_utils.py:458: TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results.
  or len(self.key_cache[layer_idx]) == 0  # the layer has no cache
/opt/app-root/src/.cache/huggingface/modules/transformers_modules/AI4Chem/CHEMLLM-2b-1_5/215c0dbc89417a06bbc3bae43a3ad61e58f0a56e/modeling_internlm2.py:1095: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if sequence_length != 1:
/opt/app-root/src/.cache/huggingface/modules/transformers_modules/AI4Chem/CHEMLLM-2b-1_5/215c0dbc89417a06bbc3bae43a3ad61e58f0a56e/modeling_internlm2.py:147: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if seq_len > self.max_position_embeddings:
/opt/app-root/lib64/python3.9/site-packages/transformers/cache_utils.py:443: TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results.
  elif len(self.key_cache[layer_idx]) == 0:  # fills previously skipped layers; checking for tensor causes errors
INFO:nncf:Statistics of the bitwidth distribution:
┍━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┑
│   Num bits (N) │ % all parameters (layers)   │ % ratio-defining parameters (layers)   │
┝━━━━━━━━━━━━━━━━┿━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┿━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┥
│              8 │ 100% (122 / 122)            │ 100% (122 / 122)                       │
┕━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┙
Applying Weight Compression ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% • 0:00:26 • 0:00:00
Model load and optimization time: 149.88100218772888 seconds

# Inference
start_time = time.time()

prompt = "Predict the aqueous solubility (in mg/mL) and explain the factors affecting solubility for: CC1=CC=C(C=C1)C(=O)CCCN"

inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)

generation_config = GenerationConfig(
    do_sample=True,
    top_k=10,
    temperature=0.9,
    max_new_tokens=500,
    repetition_penalty=1.2,
    pad_token_id=tokenizer.eos_token_id
)

# Generate output with the optimized model
outputs = model.generate(
    **inputs,
    generation_config=generation_config
)

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)

end_time = time.time()
execution_time = end_time - start_time
print(f"Inference time: {execution_time} seconds")

Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[21], line 18
      8 generation_config = GenerationConfig(
      9     do_sample=True,
     10     top_k=10,
   (...)
     14     pad_token_id=tokenizer.eos_token_id
     15 )
     17 # Generate output with the optimized model
---> 18 outputs = model.generate(
     19     **inputs,
     20     generation_config=generation_config
     21 )
     23 generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     24 print(generated_text)

File /opt/app-root/lib64/python3.9/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    112 @functools.wraps(func)
    113 def decorate_context(*args, **kwargs):
    114     with ctx_factory():
--> 115         return func(*args, **kwargs)

File /opt/app-root/lib64/python3.9/site-packages/optimum/intel/openvino/modeling_decoder.py:723, in OVModelForCausalLM.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
    721 if is_beam_search:
    722     self._first_iter_beam_search = True
--> 723 result = super().generate(
    724     inputs,
    725     generation_config,
    726     logits_processor,
    727     stopping_criteria,
    728     prefix_allowed_tokens_fn,
    729     synced_gpus,
    730     assistant_model,
    731     streamer,
    732     negative_prompt_ids,
    733     negative_prompt_attention_mask,
    734     **kwargs,
    735 )
    736 return result

File /opt/app-root/lib64/python3.9/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    112 @functools.wraps(func)
    113 def decorate_context(*args, **kwargs):
    114     with ctx_factory():
--> 115         return func(*args, **kwargs)

File /opt/app-root/lib64/python3.9/site-packages/transformers/generation/utils.py:2215, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   2207     input_ids, model_kwargs = self._expand_inputs_for_generation(
   2208         input_ids=input_ids,
   2209         expand_size=generation_config.num_return_sequences,
   2210         is_encoder_decoder=self.config.is_encoder_decoder,
   2211         **model_kwargs,
   2212     )
   2214     # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
-> 2215     result = self._sample(
   2216         input_ids,
   2217         logits_processor=prepared_logits_processor,
   2218         stopping_criteria=prepared_stopping_criteria,
   2219         generation_config=generation_config,
   2220         synced_gpus=synced_gpus,
   2221         streamer=streamer,
   2222         **model_kwargs,
   2223     )
   2225 elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH):
   2226     # 11. prepare beam search scorer
   2227     beam_scorer = BeamSearchScorer(
   2228         batch_size=batch_size,
   2229         num_beams=generation_config.num_beams,
   (...)
   2234         max_length=generation_config.max_length,
   2235     )

File /opt/app-root/lib64/python3.9/site-packages/transformers/generation/utils.py:3249, in GenerationMixin._sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, **model_kwargs)
   3247     probs = nn.functional.softmax(next_token_scores, dim=-1)
   3248     # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution
-> 3249     next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
   3250 else:
   3251     next_tokens = torch.argmax(next_token_scores, dim=-1)

File /opt/app-root/lib64/python3.9/site-packages/nncf/torch/dynamic_graph/wrappers.py:86, in wrap_operator.<locals>.wrapped(*args, **kwargs)
     84 ctx = get_current_context()
     85 if not ctx or getattr(ctx, "in_operator", False) or not ctx.is_tracing:
---> 86     op1 = operator(*args, **kwargs)
     87     return op1
     89 ctx.in_operator = True

RuntimeError: probability tensor contains either `inf`, `nan` or element < 0