Update app.py

app.py

@@ -1,4 +1,4 @@
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import hf_hub_download, snapshot_download, login
 import pandas as pd
 import importlib
 import importlib.util
@@ -6,12 +6,18 @@ import streamlit as st
 import sys
 import os
 from pathlib import Path
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+
 
-from huggingface_hub import login
 #login(token=os.environ.get("HF_TOKEN_LLAMA"))
 HF_TOKEN = os.environ.get("HF_TOKEN") #get HF_TOKEN 
 login(token=HF_TOKEN)
 
+# Check if the token is valid and has access
+print(whoami())
+
 #USER_NAME = os.getenv("USER_NAME", "").strip().strip('"')
 #PRIVATE_SPACE_NAME = os.getenv("PRIVATE_SPACE_NAME", "").strip().strip('"')
 
@@ -44,6 +50,17 @@ REPO_TYPE = "space"
 # download_private_assets(cache_dir)
 
 #Download the entire space (optional, if needed)
+# repo_dir = snapshot_download(
+#     repo_id=REPO_ID,
+#     repo_type=REPO_TYPE,
+#     token=HF_TOKEN,
+#     cache_dir="private_space_cache"
+# )
+
+# # Add repo directory to sys.path so Python can find modules inside it
+# sys.path.append(repo_dir)
+
+# Download the entire space, including the fine-tuned model folder
 repo_dir = snapshot_download(
     repo_id=REPO_ID,
     repo_type=REPO_TYPE,
@@ -51,9 +68,34 @@ repo_dir = snapshot_download(
     cache_dir="private_space_cache"
 )
 
-# Add repo directory to sys.path so Python can find modules inside it
+# The fine-tuned model is located in "fine_tuned_llama3" inside the Space
+fine_tuned_model_path = os.path.join(repo_dir, "fine_tuned_llama3")
+
+# Verify the model path
+if not os.path.exists(fine_tuned_model_path):
+    raise FileNotFoundError(f"Fine-tuned model not found at {fine_tuned_model_path}")
+
+# Add repo directory to sys.path
 sys.path.append(repo_dir)
 
+# Load the base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    "meta-llama/Meta-Llama-3-8B",
+    token=HF_TOKEN
+)
+
+# Load fine-tuned adapter (PEFT)
+fine_tuned_model = PeftModel.from_pretrained(base_model, fine_tuned_model_path)
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(fine_tuned_model_path)
+
+print("Fine-tuned model loaded successfully!")
+
+
+
+
+
 
 
 # Download specific files (if snapshot_download wasn't used)