Mediocreatmybest commited on
Commit
4fa90b9
1 Parent(s): 765c96b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -1,6 +1,6 @@
1
  import torch
2
  import gradio as gr
3
- from transformers import pipeline
4
  import ast
5
 
6
  translation_task_names = {
@@ -58,6 +58,7 @@ def translate_text(model_choice, task_choice, text_input, load_in_8bit, device):
58
  device=device, # Use selected device
59
  model_kwargs=model_kwargs,
60
  torch_dtype=dtype, # Set the floating point
 
61
  use_fast=True
62
  )
63
  # Store the loaded model
 
1
  import torch
2
  import gradio as gr
3
+ from optimum.pipelines import pipeline
4
  import ast
5
 
6
  translation_task_names = {
 
58
  device=device, # Use selected device
59
  model_kwargs=model_kwargs,
60
  torch_dtype=dtype, # Set the floating point
61
+ accelerator="bettertransformer", # Use optimum bettertransformer
62
  use_fast=True
63
  )
64
  # Store the loaded model