tosin2013 commited on
Commit
fbe81da
1 Parent(s): 194dd4b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -2
app.py CHANGED
@@ -18,11 +18,14 @@ models = {
18
 
19
  loaded_models = {}
20
 
 
 
 
21
  def load_model(model_name):
22
  if model_name not in loaded_models:
23
- loaded_models[model_name] = pipeline("question-answering", model=models[model_name])
24
  return loaded_models[model_name]
25
-
26
  def generate_score_chart(score):
27
  plt.figure(figsize=(6, 4))
28
  plt.bar(["Confidence Score"], [score], color='skyblue')
 
18
 
19
  loaded_models = {}
20
 
21
+ # Ensure we're using the CPU if GPU isn't available or necessary
22
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
23
+
24
  def load_model(model_name):
25
  if model_name not in loaded_models:
26
+ loaded_models[model_name] = pipeline("question-answering", model=models[model_name], device=0 if torch.cuda.is_available() else -1)
27
  return loaded_models[model_name]
28
+
29
  def generate_score_chart(score):
30
  plt.figure(figsize=(6, 4))
31
  plt.bar(["Confidence Score"], [score], color='skyblue')