update model naming

app.py

@@ -1,6 +1,6 @@
 # install custom package
 import os
-os.system("pip install --no-deps ./mosaicfm-0.1.2-py3-none-any.whl")
+os.system("pip install --no-deps ./tahoex-0.1.2-py3-none-any.whl")
 
 # imports
 import gc
@@ -16,9 +16,9 @@ import scanpy as sc
 from pathlib import Path
 from composer import Trainer, Callback
 from omegaconf import OmegaConf as om
-from mosaicfm.model.model import ComposerSCGPTModel
-from mosaicfm.data import CountDataset, DataCollator
-from mosaicfm.tokenizer import GeneVocab
+from tahoex.model.model import ComposerTX
+from tahoex.data import CountDataset, DataCollator
+from tahoex.tokenizer import GeneVocab
 
 # hardcoded configuration
 EMB_KEY = "X_tx1-70m"
@@ -109,7 +109,7 @@ class GradioProgressCallback(Callback):
     def predict_batch_end(self, state, logger):
         self.seen += 1
         frac = self.start + (self.end - self.start) * (self.seen / self.total)
-        self.progress(frac, desc=f"computing Tx1 embeddings ({self.seen} / {self.total} batches)...")
+        self.progress(frac, desc=f"computing Tx1 embeddings ({self.seen} / {self.total} batches)")
 
 # compute embeddings
 def _embed(adata_bytes, layer_name, feature_col, progress):
@@ -158,7 +158,7 @@ def _embed(adata_bytes, layer_name, feature_col, progress):
     vocab = GeneVocab.from_file(vocab_path)
 
     # load model
-    model = ComposerSCGPTModel(model_config=model_config, collator_config=collator_config)
+    model = ComposerTX(model_config=model_config, collator_config=collator_config)
     model.load_state_dict(torch.load(ckpt)["state"]["model"], strict=strict)
     model.to("cuda")
     model.eval()
@@ -276,12 +276,12 @@ def run_pipeline(fileobj, layer_choice, var_choice, obs_choice, progress=gr.Prog
         raise gr.Error("Please select a .var column.")
 
     # read upload file to bytes so the GPU function can load it
-    progress(0.05, desc="reading AnnData...")
+    progress(0.05, desc="reading AnnData")
     with open(fileobj.name, "rb") as f:
         adata_bytes = f.read()
 
     # compute embeddings on GPU
-    progress(0.10, desc="computing Tx1 embeddings...")
+    progress(0.10, desc="computing Tx1 embeddings")
     E, layers, var_cols, obs_cols, adata_with_emb_bytes = _embed(
         adata_bytes=adata_bytes,
         layer_name=(None if layer_choice in [None, "", "<use .X>"] else layer_choice),
@@ -297,13 +297,13 @@ def run_pipeline(fileobj, layer_choice, var_choice, obs_choice, progress=gr.Prog
         adata = sc.read_h5ad(tmp_in, backed=None)
 
     # compute UMAP
-    progress(0.60, desc="computing UMAP...")
+    progress(0.60, desc="computing UMAP")
     color_series = adata.obs[obs_choice] if (obs_choice and obs_choice in adata.obs) else None
     coords = _compute_umap_from_emb(E)
     adata.obsm["X_umap"] = coords
 
     # plot UMAP
-    progress(0.80, desc="plotting UMAP...")
+    progress(0.80, desc="plotting UMAP")
     import matplotlib.pyplot as plt
     fig = plt.figure(figsize=(5.5, 5.0))
     ax = fig.add_subplot(111)
@@ -331,7 +331,7 @@ def run_pipeline(fileobj, layer_choice, var_choice, obs_choice, progress=gr.Prog
     plt.close(fig)
 
     # save other outputs and return paths
-    progress(0.95, desc="saving outputs...")
+    progress(0.95, desc="saving outputs")
     parquet_path, h5ad_path = _save_outputs(adata, E)
     progress(1.00, desc="finished!")
     return str(umap_png.resolve()), str(parquet_path.resolve()), str(h5ad_path.resolve()), ["<use .X>"] + layers, var_cols, obs_cols