Upload app.py with huggingface_hub

app.py

@@ -21,6 +21,13 @@ SENTIMENT_CLASSES = ["Positive", "Negative", "Neutral"]
 
 print(f"===== Startup at {datetime.datetime.now()} =====")
 
+bnb = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True
+)
+
 print("Loading tokenizer...")
 tokenizer = PreTrainedTokenizerFast.from_pretrained(TOK_MODEL, token=HF_TOKEN)
 tokenizer.pad_token    = tokenizer.eos_token
@@ -28,43 +35,30 @@ tokenizer.pad_token_id = 0
 tokenizer.padding_side = "right"
 print("Tokenizer ready")
 
-print("Loading base model in 4-bit...")
-bnb = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True
-)
-base = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    quantization_config=bnb,
-    device_map="cpu",
-    token=HF_TOKEN
-)
-print("Base model ready")
-
-print("Loading classifier heads...")
-
-news_model = PeftModel.from_pretrained(base, NEWS_ADAPTER, token=HF_TOKEN)
-news_model.eval()
-news_head_path = hf_hub_download(repo_id=NEWS_ADAPTER, filename="classifier_head.pt", token=HF_TOKEN)
-news_head = nn.Linear(4096, len(NEWS_CLASSES))
-news_head.load_state_dict(torch.load(news_head_path, map_location="cpu"))
-news_head.eval()
-
-writing_model = PeftModel.from_pretrained(base, WRITING_ADAPTER, token=HF_TOKEN)
-writing_model.eval()
-writing_head_path = hf_hub_download(repo_id=WRITING_ADAPTER, filename="classifier_head.pt", token=HF_TOKEN)
-writing_head = nn.Linear(4096, len(WRITING_CLASSES))
-writing_head.load_state_dict(torch.load(writing_head_path, map_location="cpu"))
-writing_head.eval()
-
-sentiment_model = PeftModel.from_pretrained(base, SENTIMENT_ADAPTER, token=HF_TOKEN)
-sentiment_model.eval()
-sentiment_head_path = hf_hub_download(repo_id=SENTIMENT_ADAPTER, filename="classifier_head.pt", token=HF_TOKEN)
-sentiment_head = nn.Linear(4096, len(SENTIMENT_CLASSES))
-sentiment_head.load_state_dict(torch.load(sentiment_head_path, map_location="cpu"))
-sentiment_head.eval()
+def load_model_and_head(adapter_repo, num_classes):
+    print(f"Loading {adapter_repo}...")
+    base = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        quantization_config=bnb,
+        device_map="cpu",
+        token=HF_TOKEN
+    )
+    model = PeftModel.from_pretrained(base, adapter_repo, token=HF_TOKEN)
+    model.eval()
+    head_path = hf_hub_download(
+        repo_id=adapter_repo,
+        filename="classifier_head.pt",
+        token=HF_TOKEN
+    )
+    head = nn.Linear(4096, num_classes)
+    head.load_state_dict(torch.load(head_path, map_location="cpu"))
+    head.eval()
+    print(f"{adapter_repo} ready")
+    return model, head
+
+news_model,      news_head      = load_model_and_head(NEWS_ADAPTER,      5)
+writing_model,   writing_head   = load_model_and_head(WRITING_ADAPTER,   4)
+sentiment_model, sentiment_head = load_model_and_head(SENTIMENT_ADAPTER, 3)
 
 print("All models ready!")