sohomghosh commited on
Commit
381e102
1 Parent(s): 11a27f6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -1,5 +1,6 @@
1
  import pickle
2
  import gradio as gr
 
3
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
4
  from transformers import BertTokenizer, BertForSequenceClassification, pipeline, AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline, AutoModelForSeq2SeqLM, AutoModel, RobertaModel, RobertaTokenizer
5
  from sentence_transformers import SentenceTransformer
@@ -10,7 +11,7 @@ from fin_readability_sustainability import BERTClass, do_predict
10
  tokenizer_read = BertTokenizer.from_pretrained('ProsusAI/finbert')
11
 
12
 
13
-
14
  model_read = BERTClass(2, "readability")
15
  model_read.to(device)
16
  model_read.load_state_dict(torch.load('readability_model.bin', map_location=device)['model_state_dict'])
 
1
  import pickle
2
  import gradio as gr
3
+ import torch
4
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
5
  from transformers import BertTokenizer, BertForSequenceClassification, pipeline, AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline, AutoModelForSeq2SeqLM, AutoModel, RobertaModel, RobertaTokenizer
6
  from sentence_transformers import SentenceTransformer
 
11
  tokenizer_read = BertTokenizer.from_pretrained('ProsusAI/finbert')
12
 
13
 
14
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
  model_read = BERTClass(2, "readability")
16
  model_read.to(device)
17
  model_read.load_state_dict(torch.load('readability_model.bin', map_location=device)['model_state_dict'])