Update app.py

app.py

@@ -5,6 +5,7 @@ os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
 import gradio as gr
 import tempfile
 import torch
+import requests
 
 from huggingface_hub import HfApi, ModelCard, whoami
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
@@ -74,17 +75,32 @@ def run_command(command):
 
 ###########
 
+def guess_base_model(ft_model_id):
+    res = requests.get(f"https://huggingface.co/api/models/{ft_model_id}")
+    res = res.json()
+    for tag in res["tags"]:
+        if tag.startswith("base_model:"):
+            return tag.split(":")[-1]
+    raise Exception("Cannot guess the base model, please enter it manually")
+
 
 def process_model(ft_model_id: str, base_model_id: str, rank: str, private_repo, oauth_token: gr.OAuthToken | None):
     if oauth_token is None or oauth_token.token is None:
         raise gr.Error("You must be logged in")
     model_name = ft_model_id.split('/')[-1]
 
+    # validate the oauth token
+    whoami(oauth_token.token)
+
     if not os.path.exists("outputs"):
         os.makedirs("outputs")
 
     try:
         api = HfApi(token=oauth_token.token)
+
+        if not base_model_id:
+            base_model_id = guess_base_model(ft_model_id)
+            print("guess_base_model", base_model_id)
         
         with tempfile.TemporaryDirectory(dir="outputs") as outputdir:
             device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -102,7 +118,7 @@ def process_model(ft_model_id: str, base_model_id: str, rank: str, private_repo,
             print("output_stdout", output_stdout)
             print("output_stderr", output_stderr)
             if returncode != 0:
-                raise Exception(f"Error converting to LoRA PEFT {q_method}: {output_stderr}")
+                raise Exception(f"Error converting to LoRA PEFT {output_stderr}")
             print("Model converted to LoRA PEFT successfully!")
             print(f"Converted model path: {outputdir}")
 
@@ -146,8 +162,8 @@ with gr.Blocks(css=css) as demo:
     )
 
     base_model_id = HuggingfaceHubSearch(
-        label="Base model repository",
-        placeholder="Base model",
+        label="Base model repository (optional)",
+        placeholder="If empty, it will be guessed from repo tags",
         search_type="model",
     )