Update app.py

app.py

@@ -1,12 +1,16 @@
 import gradio as gr
 from transformers import AutoTokenizer
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, login
 
 api = HfApi()
 
 # Define a function to calculate tokens
-def count_tokens(llm_name, input_text):
+def count_tokens(llm_name, input_text, api_token):
     try:
+        # Login using the API token if provided
+        if api_token:
+            login(api_token)
+        
         # Load the tokenizer for the selected transformer-based model
         tokenizer = AutoTokenizer.from_pretrained(llm_name)
         tokens = tokenizer.encode(input_text)
@@ -14,16 +18,26 @@ def count_tokens(llm_name, input_text):
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Fetch model details including metadata (like download count, likes)
+# Fetch model details including metadata (like tags)
 models = list(api.list_models(task="text-generation"))
 
-# Sort models by number of downloads or likes
-llm_options = [model.modelId for model in models]#sorted(model_details, key=lambda x: x['downloads'], reverse=True)
+# Filter models that have the 'text-generation-inference' tag and 'text-generation' pipeline_tag
+filtered_models = []
+for model in models:
+    model_info = api.model_info(model.modelId)
+    if 'text-generation-inference' in model_info.tags and model_info.pipeline_tag == 'text-generation':
+        filtered_models.append(model.modelId)
 
-# Define custom CSS for a bluish theme
+# Define custom CSS for a bluish theme and cursor pointer
 custom_css = """
+.gr-dropdown {
+    cursor: pointer;
+}
 """
 
+# Set the default model to the first filtered model, or "gpt2" if there are no filtered models
+default_model = filtered_models[0] if filtered_models else "gpt2"
+
 # Create the Gradio interface
 with gr.Blocks(css=custom_css) as demo:
     gr.HTML("<h1 style='text-align: center; color: #0078d7;'>Token Counter for Transformer-Based Models</h1>")
@@ -32,14 +46,16 @@ with gr.Blocks(css=custom_css) as demo:
         "using selected transformer-based models from Hugging Face."
     )
     with gr.Row():
-        llm_dropdown = gr.Dropdown(choices=llm_options, label="Select Transformer Model", value="gpt2")
+        llm_dropdown = gr.Dropdown(choices=filtered_models, label="Select Transformer Model", value=default_model)
     with gr.Row():
         input_text = gr.Textbox(label="Enter your text")
         output = gr.Textbox(label="Token Count", interactive=False)
+    with gr.Row():
+        api_token_input = gr.Textbox(label="Enter Hugging Face API Token (if needed)", type="password", placeholder="Your API Token", interactive=True)
     with gr.Row():
         submit_btn = gr.Button("Calculate Tokens")
 
-    submit_btn.click(count_tokens, inputs=[llm_dropdown, input_text], outputs=output)
+    submit_btn.click(count_tokens, inputs=[llm_dropdown, input_text, api_token_input], outputs=output)
 
 # Launch the app
-demo.launch(share=True)
+demo.launch(share=True, debug=True)