Update app.py

app.py

@@ -1,52 +1,75 @@
+
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from transformers import pipeline
 
-# Define your actions
-ACTIONS = [
-    "move_towards_player",
-    "stand_still",
-    "follow_player",
-    "run_away",
-    "jump",
-    "mine_block",
-    "place_block",
-    "attack",
-    "use_item",
-    "chat_only"
-]
+# Define your Minecraft actions with descriptions
+ACTIONS = {
+    "move_towards_player": "walk or move toward a specific player",
+    "stand_still": "remain in place, don't move",
+    "follow_player": "continuously follow behind a player",
+    "run_away": "move away from a player or danger",
+    "jump": "jump up or over something",
+    "mine_block": "break or mine a block",
+    "place_block": "place or put down a block",
+    "attack": "attack or fight something",
+    "use_item": "use or activate an item",
+    "chat_only": "just chat without taking physical action"
+}
 
-# Load model and tokenizer - CORRECTED MODEL NAME
-model_name = "microsoft/deberta-v3-small"  # Correct format for model name
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(ACTIONS))
+# Load the zero-shot classification pipeline
+# Using a model that works well with Spaces and has good zero-shot performance
+classifier = pipeline("zero-shot-classification", 
+                      model="facebook/bart-large-mnli")
 
-# This function will need to be trained with your data
 def predict_action(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    outputs = model(**inputs)
-    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    predicted_class = torch.argmax(probabilities, dim=-1).item()
+    # Extract just the action names for classification
+    action_names = list(ACTIONS.keys())
+    
+    # Add context to improve classification by combining action names with descriptions
+    candidate_labels = [f"{action}: {ACTIONS[action]}" for action in action_names]
+    
+    # Run zero-shot classification
+    result = classifier(text, candidate_labels, multi_label=False)
+    
+    # Extract results
+    action_scores = {}
+    for i, label in enumerate(result["labels"]):
+        # Extract just the action name from the combined label+description
+        action_name = label.split(":")[0]
+        action_scores[action_name] = result["scores"][i]
+    
+    # Find the best action
+    best_action = max(action_scores, key=action_scores.get)
+    confidence = action_scores[best_action]
     
-    # Return both the action and the confidence score
     return {
-        "action": ACTIONS[predicted_class],
-        "confidence": probabilities[0][predicted_class].item(),
-        "all_actions": {ACTIONS[i]: probabilities[0][i].item() for i in range(len(ACTIONS))}
+        "action": best_action,
+        "confidence": confidence,
+        "all_actions": action_scores
     }
 
-# Create the Gradio interface
-def process_text(text):
-    result = predict_action(text)
-    return result["action"], result["confidence"], result["all_actions"]
+# Example to test (you can remove this in production)
+example_text = "I want to go and meet that player over there"
+test_result = predict_action(example_text)
+print(f"Example: '{example_text}'")
+print(f"Best action: {test_result['action']}")
+print(f"Confidence: {test_result['confidence']:.4f}")
 
-# Define the Gradio interface
+# Create the Gradio interface
 demo = gr.Interface(
     fn=predict_action,
-    inputs=gr.Textbox(label="Character text"),
-    outputs=gr.JSON(label="Result"),
-    title="Minecraft Action Predictor",
-    description="Predicts the best action based on character text"
+    inputs=gr.Textbox(label="Character text from C.AI"),
+    outputs=gr.JSON(label="Predicted Action"),
+    title="Minecraft NPC Action Predictor",
+    description="""This tool analyzes what your C.AI character says and determines the best action to take in Minecraft.
+                   Input the character's dialogue to get the appropriate in-game action.""",
+    examples=[
+        ["I'll help you build that house!"],
+        ["Look out, there's a creeper behind you!"],
+        ["Let me show you the way to the village."],
+        ["I think I'll just wait here for a while."],
+        ["I need to break this stone to make a path."]
+    ]
 )
 
 demo.launch()