Update app.py

app.py

@@ -6,40 +6,59 @@ import gradio as gr
 from sklearn.cluster import AgglomerativeClustering
 import re
 import pandas as pd
+from keybert.backend import TransformerBackend
 
 
 def load_models():
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    pubmed_model = AutoModel.from_pretrained("microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract",
-                                             use_safetensors=True).to(device)
+
+    pubmed_tok = AutoTokenizer.from_pretrained(
+        "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract"
+    )
+    pubmed_model = AutoModel.from_pretrained(
+        "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract",
+        use_safetensors=True
+    ).to(device)
+
     try:
         nlp = spacy.load("en_core_sci_sm")
     except OSError:
-        import os
-        os.system("python -m spacy download en_core_sci_sm")
+        url = ("https://huggingface.co/allenai/scispacy_models/resolve/main/"
+               "en_core_sci_sm-0.5.5.tar.gz")
+        os.system(f"{sys.executable} -m pip install {url}")
         nlp = spacy.load("en_core_sci_sm")
 
+    led_tok = LEDTokenizer.from_pretrained("dancessa/led_pubmed_summarization")
+    led_mod = LEDForConditionalGeneration.from_pretrained(
+        "dancessa/led_pubmed_summarization",
+        use_safetensors=True
+    ).to(device)
+
+    kb_backend = TransformerBackend(pubmed_model, pubmed_tok)
+    kb_model   = KeyBERT(model=kb_backend)
+
     return {
         "led": {
-            "tokenizer": LEDTokenizer.from_pretrained("dancessa/led_pubmed_summarization"),
-            "model": LEDForConditionalGeneration.from_pretrained("dancessa/led_pubmed_summarization",
-                                                                 use_safetensors=True).to(device)
+            "tokenizer": led_tok,
+            "model": led_mod
         },
         "pubmed": {
-            "tokenizer": AutoTokenizer.from_pretrained("microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract"),
+            "tokenizer": pubmed_tok,
             "model": pubmed_model
         },
-        "keybert": KeyBERT(model=pubmed_model),
+        "keybert": kb_model,
         "spacy": nlp
     }
-
+    
 
 models = load_models()
-nlp = models["spacy"]
-kw_model = models["keybert"]
-pubmed_tokenizer = models["pubmed"]["tokenizer"]
-led_model = models["led"]["model"]
-led_tokenizer = models["led"]["tokenizer"]
+nlp             = models["spacy"]
+kw_model        = models["keybert"]
+pubmed_tokenizer= models["pubmed"]["tokenizer"]
+pubmed_model    = models["pubmed"]["model"]
+led_tokenizer   = models["led"]["tokenizer"]
+led_model       = models["led"]["model"]
+
 
 
 def generate_summary(medical_text):