Jyotiyadav commited on
Commit
4282a0d
·
verified ·
1 Parent(s): 39fb2f8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -2,14 +2,15 @@ import gradio as gr
2
  import torch
3
  import torchvision
4
  from transformers import pipeline
 
5
 
6
  # Initialize the pipeline
7
  pipe = pipeline(
8
  "text-generation",
9
  model="DataIntelligenceTeam/mistral_7B_NER",
10
  torch_dtype=torch.bfloat16,
11
- device_map="auto"
12
- )
13
 
14
  # Function to generate output based on input
15
  def generate_output(input_text):
 
2
  import torch
3
  import torchvision
4
  from transformers import pipeline
5
+ auth_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
6
 
7
  # Initialize the pipeline
8
  pipe = pipeline(
9
  "text-generation",
10
  model="DataIntelligenceTeam/mistral_7B_NER",
11
  torch_dtype=torch.bfloat16,
12
+ device_map="auto",
13
+ use_auth_token=auth_token)
14
 
15
  # Function to generate output based on input
16
  def generate_output(input_text):