Update app.py

app.py

@@ -3,34 +3,36 @@ import os
 from huggingface_hub import HfApi, login
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-def process(model_id, dataset):
-    print("111")
-    # Download Sample Model from Hugging Face to Publish Again 
+# NVidia A10G Large
+
+# google/gemma-2-9b-it
+# meta-llama/Meta-Llama-3-8B-Instruct
+
+profile = "bstraehle"
+
+def download_tokenizer_and_model(model_id)
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     model = AutoModelForCausalLM.from_pretrained(model_id)
-    # Local Path of Model
-    print("222")
-    model_path = model_id
-    model.save_pretrained(model_path)
+    model.save_pretrained(model_id)
+    return tokenizer, model
+
+def upload_model(tokenizer, model)
+    model_name = model_id[model_id.rfind('/')+1:]
+    print(model_name)
+    model_repo_name = f"{profile}/{model_name}"
     login(token=os.environ["HF_TOKEN"])
     api = HfApi()
-    model_repo_name = "bstraehle/Meta-Llama-3-8B-Instruct"
-    
-    #Create Repo in Hugging Face
-    print("333")
     api.create_repo(repo_id=model_repo_name)
-    
-    #Upload Model folder from Local to HuggingFace 
-    print("444")
     api.upload_folder(
         folder_path=model_path,
         repo_id=model_repo_name
     )
+    tokenizer.push_to_hub(model_repo_name)
+    
+def process(model_id, dataset):
+    tokenizer, model = download_model(model_id)
+    upload_model(tokenizer, model)
     
-    # Publish Model Tokenizer on Hugging Face
-    print("555")
-    tokenizer.push_to_hub(model_repo_name)    
-
     return "Done"
 
 demo = gr.Interface(fn=process,