add great readme

app.py

@@ -1,15 +1,13 @@
 import gradio as gr
 import torch
-from transformers import AutoModel, BitsAndBytesConfig
+from transformers import AutoModel, BitsAndBytesConfig, AutoTokenizer
 import tempfile
 from huggingface_hub import HfApi
 from huggingface_hub import list_models
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from bitsandbytes.nn import Linear4bit
-from packaging import version
 import os
-from tqdm import tqdm
-
+from huggingface_hub import snapshot_download
 
 def hello(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> str:
     # ^ expect a gr.OAuthProfile object as input to get the user's profile
@@ -42,11 +40,52 @@ def check_model_exists(
 def create_model_card(
     model_name, quant_type_4, double_quant_4, compute_type_4, quant_storage_4
 ):
-    model_card = f"""---
+    # Try to download the original README
+    original_readme = ""
+    original_yaml_header = ""
+    try:
+        # Download the README.md file from the original model
+        model_path = snapshot_download(repo_id=model_name, allow_patterns=["README.md"], repo_type="model")
+        readme_path = os.path.join(model_path, "README.md")
+        
+        if os.path.exists(readme_path):
+            with open(readme_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+                
+                if content.startswith('---'):
+                    parts = content.split('---', 2)
+                    if len(parts) >= 3:
+                        original_yaml_header = parts[1]
+                        original_readme = '---'.join(parts[2:])
+                    else:
+                        original_readme = content
+                else:
+                    original_readme = content
+    except Exception as e:
+        print(f"Error reading original README: {str(e)}")
+        original_readme = ""
+    
+    # Create new YAML header with base_model field
+    yaml_header = f"""---
 base_model:
-- {model_name}
----
-
+- {model_name}"""
+    
+    # Add any original YAML fields except base_model
+    if original_yaml_header:
+        skip_next_line = False
+        for line in original_yaml_header.strip().split('\n'):
+            if skip_next_line:
+                skip_next_line = False
+                continue
+            if line.strip().startswith('base_model:'):
+                skip_next_line = True
+                continue
+            yaml_header += f"\n{line}"
+    # Complete the YAML header
+    yaml_header += "\n---"
+    
+    # Create the quantization info section
+    quant_info = f"""
 # {model_name} (Quantized)
 
 ## Description
@@ -63,6 +102,13 @@ It's quantized using the BitsAndBytes library to 4-bit using the [bnb-my-repo](h
 
 """
 
+    # Combine everything
+    model_card = yaml_header + quant_info
+    
+    # Append original README content if available
+    if original_readme and not original_readme.isspace():
+        model_card += "\n\n# Original Model Information\n" + original_readme
+    
     return model_card
 
 
@@ -138,6 +184,8 @@ def save_model(
 
     with tempfile.TemporaryDirectory() as tmpdirname:
         # Save model
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=auth_token.token)
+        tokenizer.save_pretrained(tmpdirname, safe_serialization=True, use_auth_token=auth_token.token)
         model.save_pretrained(
             tmpdirname, safe_serialization=True, use_auth_token=auth_token.token
         )