Add workaround for gated models

README.md

@@ -9,11 +9,6 @@ app_file: app.py
 pinned: false
 license: apache-2.0
 short_description: Prompt enhancing models interface
-
-hf_oauth: true
-hf_oauth_scopes:
- - read-repos
- - manage-repos
 ---
 
 A playground to test and compare several prompt enhancing models.

app.py

@@ -1,6 +1,7 @@
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import login
 
 
 choices_base_models = {
@@ -17,11 +18,18 @@ choices_gen_token = {
     'groloch/Ministral-3b-instruct-PromptEnhancing': 'ministral/Ministral-3b-instruct'
 }
 
+gated_models = [
+    'groloch/Llama-3.2-3B-Instruct-PromptEnhancing',
+    'groloch/gemma-2-2b-it-PromptEnhancing'
+]
+
 previous_choice = ''
 
 model = None
 tokenizer = None
 
+logged_in = False
+
 
 def load_model(adapter_repo_id: str):
     global model, tokenizer
@@ -37,7 +45,8 @@ def generate(prompt_to_enhance: str,
              max_tokens: float,
              temperature: float, 
              top_p: float, 
-             repetition_penalty: float
+             repetition_penalty: float,
+             access_token: str
              ):
     if prompt_to_enhance is None or prompt_to_enhance == '':
         raise gr.Error('Please enter a prompt')
@@ -47,6 +56,15 @@ def generate(prompt_to_enhance: str,
         previous_choice = choice
         load_model(choice)
         
+    if choice in gated_models and access_token == '':
+        raise gr.Error(f'Please enter your access token (in Additional inputs) if youre using one of the following \
+            models: {', '.join(gated_models)}. Make sure you have access to those models.')
+        
+    global logged_in
+    if not logged_in and choice in gated_models:
+        login(access_token)
+        logged_in = True
+        
     chat = [
         {'role' : 'user', 'content': prompt_to_enhance}
     ]
@@ -124,17 +142,23 @@ input_repetition_penalty = gr.Number(
     maximum=5.0,
     step=0.1
 )
+input_access_token = gr.Text(
+    label='Access token for gated models',
+    value=''
+)
 
 demo = gr.Interface(
     generate,
     title='Prompt Enhancing Playground',
     description='This space is a tool to compare the different prompt enhancing model I have finetuned. \
-            Feel free to experiment as you want !',
+            Feel free to experiment as you want ! \n\
+            If you want to use this locally, you can download the gpu version (see in files)',
     inputs=[input_prompt, model_choice],
     additional_inputs=[input_max_tokens, 
                        input_temperature, 
                        input_top_p, 
-                       input_repetition_penalty
+                       input_repetition_penalty,
+                       input_access_token
                        ],
     outputs=['text']
 )