yabramuvdi commited on
Commit
93a3f9a
·
verified ·
1 Parent(s): 43130a6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -40
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import os
 
2
  if os.environ.get("SPACES_ZERO_GPU") is not None:
3
  import spaces
4
  else:
@@ -11,7 +12,7 @@ else:
11
 
12
  @spaces.GPU
13
  def fake_gpu():
14
- pass
15
 
16
  import numpy as np
17
  import torch
@@ -58,7 +59,7 @@ def get_next_token_predictions(text, model_name, top_k=10):
58
 
59
  return top_k_tokens, top_k_probs.tolist()
60
 
61
- def predict_next_token(text, model_name, custom_token=""):
62
  # Add custom token if provided
63
  if custom_token:
64
  text += custom_token
@@ -69,7 +70,7 @@ def predict_next_token(text, model_name, custom_token=""):
69
  # Format predictions
70
  predictions = "\n".join([f"'{token}' : {prob:.4f}" for token, prob in zip(tokens, probs)])
71
 
72
- return text, gr.update(choices=[f"'{t}'" for t in tokens]), predictions
73
 
74
  # Create the interface
75
  with gr.Blocks() as demo:
@@ -77,31 +78,26 @@ with gr.Blocks() as demo:
77
 
78
  gr.Markdown("""
79
  This application allows you to interactively generate text using various transformer models.
80
- You can either select from the predicted next tokens or write your own tokens to continue the text generation.
81
-
82
- Select a model, start typing or choose from the predicted tokens, and see how the model continues your text!
83
  """)
84
 
85
- with gr.Row():
86
- text_input = gr.Textbox(
87
- lines=5,
88
- label="Text",
89
- placeholder="Type your text here...",
90
- value="The quick brown fox"
91
- )
92
-
93
  with gr.Row():
94
  model_dropdown = gr.Dropdown(
95
  choices=list(AVAILABLE_MODELS.keys()),
96
  value="distilgpt2",
97
  label="Select Model"
98
  )
99
-
100
  with gr.Row():
101
- custom_input = gr.Textbox(
102
- label="Custom token (optional)",
103
- placeholder="Type a custom token..."
 
 
104
  )
 
 
 
105
 
106
  with gr.Row():
107
  token_dropdown = gr.Dropdown(
@@ -115,29 +111,11 @@ with gr.Blocks() as demo:
115
  label="Token probabilities"
116
  )
117
 
118
- # Set up event handlers
119
- text_input.change(
120
- predict_next_token,
121
- inputs=[text_input, model_dropdown, custom_input],
122
- outputs=[text_input, token_dropdown, predictions_output]
123
- )
124
-
125
- model_dropdown.change(
126
- predict_next_token,
127
- inputs=[text_input, model_dropdown, custom_input],
128
- outputs=[text_input, token_dropdown, predictions_output]
129
- )
130
-
131
- custom_input.change(
132
- predict_next_token,
133
- inputs=[text_input, model_dropdown, custom_input],
134
- outputs=[text_input, token_dropdown, predictions_output]
135
- )
136
-
137
- token_dropdown.change(
138
  predict_next_token,
139
- inputs=[text_input, model_dropdown, gr.Textbox(value="")],
140
- outputs=[text_input, token_dropdown, predictions_output]
141
  )
142
 
143
  demo.queue().launch()
 
1
  import os
2
+ # Handle Spaces GPU
3
  if os.environ.get("SPACES_ZERO_GPU") is not None:
4
  import spaces
5
  else:
 
12
 
13
  @spaces.GPU
14
  def fake_gpu():
15
+ pass
16
 
17
  import numpy as np
18
  import torch
 
59
 
60
  return top_k_tokens, top_k_probs.tolist()
61
 
62
+ def predict_next_token(model_name, text, custom_token=""):
63
  # Add custom token if provided
64
  if custom_token:
65
  text += custom_token
 
70
  # Format predictions
71
  predictions = "\n".join([f"'{token}' : {prob:.4f}" for token, prob in zip(tokens, probs)])
72
 
73
+ return gr.update(choices=[f"'{t}'" for t in tokens]), predictions
74
 
75
  # Create the interface
76
  with gr.Blocks() as demo:
 
78
 
79
  gr.Markdown("""
80
  This application allows you to interactively generate text using various transformer models.
81
+ Select a model, enter your text, and click predict to see the possible next tokens and their probabilities.
 
 
82
  """)
83
 
 
 
 
 
 
 
 
 
84
  with gr.Row():
85
  model_dropdown = gr.Dropdown(
86
  choices=list(AVAILABLE_MODELS.keys()),
87
  value="distilgpt2",
88
  label="Select Model"
89
  )
90
+
91
  with gr.Row():
92
+ text_input = gr.Textbox(
93
+ lines=5,
94
+ label="Text",
95
+ placeholder="Type your text here...",
96
+ value="The quick brown fox"
97
  )
98
+
99
+ with gr.Row():
100
+ predict_button = gr.Button("Predict")
101
 
102
  with gr.Row():
103
  token_dropdown = gr.Dropdown(
 
111
  label="Token probabilities"
112
  )
113
 
114
+ # Set up predict button event handler
115
+ predict_button.click(
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
116
  predict_next_token,
117
+ inputs=[model_dropdown, text_input],
118
+ outputs=[token_dropdown, predictions_output]
119
  )
120
 
121
  demo.queue().launch()