Update app.py

app.py

@@ -10,84 +10,80 @@ AVAILABLE_MODELS = {
     "pythia-160m": "EleutherAI/pythia-160m"
 }
 
-class TextGenerator:
-    def __init__(self):
-        self.model = None
-        self.tokenizer = None
-        
-    def load_model(self, model_name: str) -> str:
-        try:
-            self.model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name])
-            self.tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name])
-            return f"Successfully loaded {model_name}"
-        except Exception as e:
-            return f"Error loading model: {str(e)}"
+generator = None
+
+def load_model(model_name):
+    global generator
+    try:
+        model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name])
+        tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name])
+        generator = (model, tokenizer)
+        return f"Successfully loaded {model_name}"
+    except Exception as e:
+        return f"Error loading model: {str(e)}"
+
+def get_predictions(text, model_name):
+    global generator
+    if not generator:
+        load_model(model_name)
     
-    def get_next_token_predictions(self, text: str, top_k: int = 10):
-        if not self.model or not self.tokenizer:
-            return [], []
-        
-        inputs = self.tokenizer(text, return_tensors="pt")
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-            logits = outputs.logits[0, -1, :]
-            probs = torch.nn.functional.softmax(logits, dim=-1)
-            
-        top_k_probs, top_k_indices = torch.topk(probs, top_k)
-        top_k_tokens = [self.tokenizer.decode([idx.item()]) for idx in top_k_indices]
+    model, tokenizer = generator
+    inputs = tokenizer(text, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits[0, -1, :]
+        probs = torch.nn.functional.softmax(logits, dim=-1)
         
-        return top_k_tokens, top_k_probs.tolist()
+    top_k_probs, top_k_indices = torch.topk(probs, k=10)
+    top_k_tokens = [tokenizer.decode([idx.item()]) for idx in top_k_indices]
+    predictions = "\n".join([f"'{token}' : {prob:.4f}" for token, prob in zip(top_k_tokens, top_k_probs)])
+    
+    return top_k_tokens, predictions
 
-generator = TextGenerator()
+def generate(model_name, text, token_choice="", custom_token=""):
+    if token_choice:
+        text += token_choice.strip("'")
+    if custom_token:
+        text += custom_token
+        
+    tokens, predictions = get_predictions(text, model_name)
+    return text, gr.Dropdown(choices=[f"'{t}'" for t in tokens]), predictions
 
-def format_predictions(tokens, probs):
-    if not tokens or not probs:
-        return "No predictions available"
+with gr.Blocks() as demo:
+    gr.Markdown("# Interactive Text Generation")
     
-    formatted = "Predicted next tokens:\n\n"
-    for token, prob in zip(tokens, probs):
-        formatted += f"'{token}' : {prob:.4f}\n"
-    return formatted
-
-def update_output(model_name, text, custom_token, selected_token):
-    output = text
+    model_name = gr.Dropdown(
+        choices=list(AVAILABLE_MODELS.keys()),
+        value="distilgpt2",
+        label="Select Model"
+    )
     
-    if not generator.model or generator.model.name_or_path != AVAILABLE_MODELS[model_name]:
-        load_message = generator.load_model(model_name)
-        if "Error" in load_message:
-            return text, "", "", gr.update(choices=[]), load_message
+    text = gr.Textbox(
+        lines=5,
+        label="Text",
+        placeholder="Type or select tokens to generate text..."
+    )
     
-    if custom_token:
-        output += custom_token
-    elif selected_token:
-        output += selected_token.strip("'")
+    with gr.Row():
+        token_choice = gr.Dropdown(
+            choices=[],
+            label="Select predicted token"
+        )
+        custom_token = gr.Textbox(
+            label="Or type custom token"
+        )
     
-    tokens, probs = generator.get_next_token_predictions(output)
-    predictions = format_predictions(tokens, probs)
-    token_choices = [f"'{token}'" for token in tokens]
+    predictions = gr.Textbox(
+        label="Predictions",
+        lines=10
+    )
     
-    return output, "", "", gr.update(choices=token_choices), predictions
-
-demo = gr.Interface(
-    fn=update_output,
-    inputs=[
-        gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="distilgpt2", label="Select Model"),
-        gr.Textbox(lines=5, label="Generated Text", placeholder="Start typing or select a token..."),
-        gr.Textbox(label="Custom Token", placeholder="Type your own token..."),
-        gr.Dropdown(choices=[], label="Select from predicted tokens")
-    ],
-    outputs=[
-        gr.Textbox(lines=5, label="Generated Text"),
-        gr.Textbox(label="Custom Token"),
-        gr.Textbox(label="Selected Token"),
-        gr.Dropdown(label="Predicted Tokens"),
-        gr.Textbox(lines=12, label="Predictions")
-    ],
-    title="Interactive Text Generation",
-    description="Generate text by selecting predicted tokens or writing your own."
-)
+    for component in [model_name, token_choice, custom_token]:
+        component.change(
+            generate,
+            inputs=[model_name, text, token_choice, custom_token],
+            outputs=[text, token_choice, predictions]
+        )
 
-if __name__ == "__main__":
-    demo.launch()
-else:
-    demo.launch(show_error=True) # Required for Hugging Face Spaces
+demo.queue().launch()