Mhammad Ibrahim commited on
Commit
311f263
·
1 Parent(s): d98367a

Fix meta tensor issue with explicit device mapping

Browse files
Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -7,12 +7,18 @@
7
  # demo.launch()
8
 
9
  import gradio as gr
 
10
  from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
11
 
12
  # Load model and tokenizer from Hugging Face Hub
13
  tokenizer = AutoTokenizer.from_pretrained("Mhammad2023/bert-finetuned-ner")
14
- model = AutoModelForTokenClassification.from_pretrained("Mhammad2023/bert-finetuned-ner", from_tf=True)
15
- classifier = pipeline("token-classification", model=model, tokenizer=tokenizer)
 
 
 
 
 
16
 
17
  def predict(text):
18
  results = classifier(text)
 
7
  # demo.launch()
8
 
9
  import gradio as gr
10
+ import torch
11
  from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
12
 
13
  # Load model and tokenizer from Hugging Face Hub
14
  tokenizer = AutoTokenizer.from_pretrained("Mhammad2023/bert-finetuned-ner")
15
+ model = AutoModelForTokenClassification.from_pretrained(
16
+ "Mhammad2023/bert-finetuned-ner",
17
+ from_tf=True,
18
+ torch_dtype=torch.float32,
19
+ device_map="cpu"
20
+ )
21
+ classifier = pipeline("token-classification", model=model, tokenizer=tokenizer, device="cpu")
22
 
23
  def predict(text):
24
  results = classifier(text)