Update app.py

app.py

@@ -1,31 +1,80 @@
 import gradio as gr
 
+# from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+# model_path = "anzorq/m2m100_418M_ft_ru-kbd_44K"  
+# src_lang="ru" 
+# tgt_lang="zu"
+
+# # tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang=src_lang)
+# tokenizer = AutoTokenizer.from_pretrained(model_path)
+# model = AutoModelForSeq2SeqLM.from_pretrained(model_path, use_safetensors=True)#, load_in_4bit=True, device_map="auto")
+# model.to_bettertransformer()
+
+# def translate(text, num_beams=4, num_return_sequences=4):
+#   inputs = tokenizer(text, return_tensors="pt")
+
+#   num_return_sequences = min(num_return_sequences, num_beams)
+
+#   translated_tokens = model.generate(
+#       **inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], num_beams=num_beams, num_return_sequences=num_return_sequences
+#   )
+
+#   translations = []
+#   for translation in tokenizer.batch_decode(translated_tokens, skip_special_tokens=True):
+#       translations.append(translation)
+
+#   # result = {"input":text, "translations":translations}
+#   return text, translations
+
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from optimum.bettertransformer import BetterTransformer
+import intel_extension_for_pytorch as ipex
+from transformers.modeling_outputs import BaseModelOutput
+import torch
 
-model_path = "anzorq/m2m100_418M_ft_ru-kbd_44K"  
-src_lang="ru" 
-tgt_lang="zu"
+model_path = "anzorq/m2m100_418M_ft_ru-kbd_44K"
+src_lang = "ru"
+tgt_lang = "zu"
 
-# tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang=src_lang)
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_path, use_safetensors=True)#, load_in_4bit=True, device_map="auto")
-model.to_bettertransformer()
+model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
 
-def translate(text, num_beams=4, num_return_sequences=4):
-  inputs = tokenizer(text, return_tensors="pt")
+# flash attention optimization
+model = BetterTransformer.transform(model, keep_original_model=False)
 
-  num_return_sequences = min(num_return_sequences, num_beams)
+# ipex optimization
+model.eval()
+model = ipex.optimize(model, dtype=torch.float, level="O1", conv_bn_folding=False, inplace=True)
 
-  translated_tokens = model.generate(
-      **inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], num_beams=num_beams, num_return_sequences=num_return_sequences
-  )
+# Get the encoder
+encoder = model.get_encoder()
 
-  translations = []
-  for translation in tokenizer.batch_decode(translated_tokens, skip_special_tokens=True):
-      translations.append(translation)
+# Prepare an example input for the encoder
+example_input_text = "Example text in Russian"
+inputs_example = tokenizer(example_input_text, return_tensors="pt")
 
-  # result = {"input":text, "translations":translations}
-  return text, translations
+# Trace just the encoder with strict=False
+scripted_encoder = torch.jit.trace(encoder, inputs_example['input_ids'], strict=False)
+
+def translate(text, num_beams=4, num_return_sequences=4):
+    inputs = tokenizer(text, return_tensors="pt")
+    num_return_sequences = min(num_return_sequences, num_beams)
+
+    # Use the scripted encoder for the first step of inference
+    encoder_output_dict = scripted_encoder(inputs['input_ids'])
+    encoder_outputs = BaseModelOutput(last_hidden_state=encoder_output_dict['last_hidden_state'])
+
+    # Use the original, untraced model for the second step, passing the encoder's outputs as inputs
+    translated_tokens = model.generate(
+        encoder_outputs=encoder_outputs,
+        forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], 
+        num_beams=num_beams, 
+        num_return_sequences=num_return_sequences
+    )
+
+    translations = [tokenizer.decode(translation, skip_special_tokens=True) for translation in translated_tokens]
+    return text, translations
 
 output = gr.Textbox()
 # with gr.Accordion("Advanced Options"):