Update app.py

app.py

@@ -207,10 +207,9 @@ def quantize(model_path, repo_id, quant_method=None):
     
     return final_path
 
-def create_readme(repo_name, base_model_name, lora_model_name, quant_methods, username):
+def create_readme(repo_name, base_model_name, lora_model_name, quant_methods):
     readme_path = os.path.join("output", repo_name, "README.md")
-    readme_template = """---
-tags:
+    readme_template = """---tags:
 - autotrain
 - text-generation-inference
 - text-generation
@@ -236,32 +235,6 @@ datasets:
 - created_at: {created_at}
 - created_by: [Steven10429/apply_lora_and_quantize](https://github.com/Steven10429/apply_lora_and_quantize)
 
-## Usage:
-```python
-
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-model_path = "{username}/{repo_name}"
-
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(
-    model_path,
-    device_map="auto",
-    torch_dtype='auto'
-).eval()
-
-# Prompt content: "hi"
-messages = [
-    {"role": "user", "content": "hi"}
-]
-
-input_ids = tokenizer.apply_chat_template(conversation=messages, tokenize=True, add_generation_prompt=True, return_tensors='pt')
-output_ids = model.generate(input_ids.to('cuda'))
-response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
-
-# Model response: "Hello! How can I assist you today?"
-print(response)
-```
 """.format(
         quantization="\n- quantization" if len(quant_methods) > 0 else "",
         base_model_name=base_model_name,
@@ -269,7 +242,6 @@ print(response)
         repo_name=repo_name,
         quant_methods=quant_methods,
         created_at=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
-        username=username
     )
     
     with open(readme_path, "w") as f:
@@ -288,8 +260,11 @@ def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf
     """
     try:
         current_logs.clear()
+        if hf_token.strip().lower() == "auto":
+            hf_token = os.getenv("HF_TOKEN")
+        elif hf_token.startswith("hf_"):
+            os.environ["HF_TOKEN"] = hf_token
         login(hf_token)
-        os.environ["HF_TOKEN"] = hf_token
         api = HfApi(token=hf_token)
         username = api.whoami()["name"]
         
@@ -312,7 +287,7 @@ def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf
         for quant_method in quant_methods:
             quantize(output_dir, repo_name, quant_method=quant_method)
             
-        create_readme(repo_name, base_model_name, lora_model_name, quant_methods, username)
+        create_readme(repo_name, base_model_name, lora_model_name, quant_methods)
         
         # 上传合并后的模型和量化模型
         api.upload_large_folder(
@@ -332,6 +307,7 @@ def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf
     except Exception as e:
         error_message = f"Error during processing: {e}"
         log(error_message)
+        raise e
         return "\n".join(current_logs)
 
 @timeit
@@ -368,7 +344,7 @@ def create_ui():
                 hf_token = gr.Textbox(
                     label="Hugging Face Token",
                     placeholder="Enter your Hugging Face Token",
-                    value=os.getenv("HF_TOKEN")
+                    value="Auto"
                 )
                 convert_btn = gr.Button("Start Conversion", variant="primary")
             with gr.Column():