Update app.py

app.py

@@ -1,93 +1,90 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
+import re
 from typing import List, Dict
 from datetime import datetime
 
-class MissionContext:
+class MissionGenerator:
     def __init__(self):
         self.mission_counter = 1
-        self.current_objectives = {}
-        self.conversation_history = []
-        
-    def add_to_history(self, role: str, content: str):
-        self.conversation_history.append({
-            "role": role,
-            "content": content,
-            "timestamp": datetime.now().isoformat()
-        })
-        # Keep only last 5 messages for context
-        if len(self.conversation_history) > 5:
-            self.conversation_history.pop(0)
+        self.keywords = {
+            'main': [
+                'must', 'need to', 'primary', 'main goal', 'objective is to',
+                'task is to', 'have to', 'required to'
+            ],
+            'secondary': [
+                'optional', 'can also', 'bonus', 'additional', 'extra',
+                'if possible', 'could also', 'might also'
+            ],
+            'alternative': [
+                'alternatively', 'another way', 'or', 'instead',
+                'different approach', 'other option', 'different way'
+            ]
+        }
 
-class MissionGenerator:
-    def __init__(self):
-        # Using FLAN-T5-base, a free and lightweight model good for instruction following
-        self.model_name = "google/flan-t5-base"
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
-        self.context = MissionContext()
+    def extract_objectives(self, text: str) -> List[Dict]:
+        sentences = [s.strip() for s in re.split(r'[.!?]+', text) if s.strip()]
+        objectives = []
         
-    def format_conversation_history(self) -> str:
-        """Format conversation history for the model input"""
-        formatted = ""
-        for msg in self.context.conversation_history:
-            role = "User" if msg["role"] == "user" else "Assistant"
-            formatted += f"{role}: {msg['content']}\n"
-        return formatted
+        for sentence in sentences:
+            sentence_lower = sentence.lower()
+            
+            # Determine objective type based on keywords and sentence position
+            obj_type = 'Main'  # Default to Main for first sentence if no other indicators
+            
+            if any(kw in sentence_lower for kw in self.keywords['alternative']):
+                obj_type = 'Alternative'
+            elif any(kw in sentence_lower for kw in self.keywords['secondary']):
+                obj_type = 'Secondary'
+            elif any(kw in sentence_lower for kw in self.keywords['main']) or not objectives:
+                obj_type = 'Main'
+                
+            # Clean up the sentence
+            clean_sentence = sentence
+            for keyword_list in self.keywords.values():
+                for kw in keyword_list:
+                    clean_sentence = re.sub(rf'\b{kw}\b', '', clean_sentence, flags=re.IGNORECASE)
+            
+            # Further cleanup
+            clean_sentence = re.sub(r'^[\s,]+|[\s,]+$', '', clean_sentence)
+            clean_sentence = clean_sentence.strip()
+            
+            if clean_sentence:
+                # Create identifier from key words
+                words = re.findall(r'\b\w+\b', clean_sentence.lower())
+                identifier = '_'.join(words[:2])
+                
+                objectives.append({
+                    'type': obj_type,
+                    'identifier': identifier,
+                    'description': clean_sentence
+                })
         
+        return objectives
+
     def generate_response(self, user_input: str) -> tuple[str, str]:
-        """Generate both conversational response and formatted mission objectives"""
-        self.context.add_to_history("user", user_input)
+        objectives = self.extract_objectives(user_input)
         
-        # Create prompt for the model
-        conversation_history = self.format_conversation_history()
-        prompt = f"""
-        Previous conversation:
-        {conversation_history}
-
-        Task: Generate a mission for Original War game based on the conversation.
-        Format the response as follows:
-        1. A conversational response understanding the mission
-        2. The mission objectives in Original War format using Add Main/Secondary/Alternative
+        # Generate formatted output
+        output = f"# M{self.mission_counter}\n"
         
-        Current request: {user_input}
-        """
+        for obj in objectives:
+            output += f"Add {obj['type']} {obj['identifier']}\n"
+            output += f"- {obj['description']}\n"
         
-        # Generate response using the model
-        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
-        outputs = self.model.generate(
-            inputs["input_ids"],
-            max_length=256,
-            do_sample=True,  # Enable sampling
-            temperature=0.7,  # Control randomness
-            top_p=0.9,       # Nucleus sampling
-            num_beams=4,     # Beam search
-            no_repeat_ngram_size=2
-        )
+        # Check for silent/suppress keywords
+        if any(word in user_input.lower() for word in ['silently', 'quietly', 'without notification', "don't notify"]):
+            output += "Dnd\n"
         
-        full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        output += "#\n"
         
-        # Split the response into conversational and formatted parts
-        try:
-            parts = full_response.split("# M1")
-            chat_response = parts[0].strip()
-            formatted_objectives = "# M1" + parts[1] if len(parts) > 1 else self.generate_fallback_objectives(user_input)
-        except Exception:
-            chat_response = full_response
-            formatted_objectives = self.generate_fallback_objectives(user_input)
+        # Generate chat response
+        chat_response = f"I've created {len(objectives)} objectives based on your description. "
+        if objectives:
+            types = [obj['type'] for obj in objectives]
+            chat_response += f"This includes {', '.join(f'{t.lower()} objectives' for t in set(types))}. "
+        chat_response += "Would you like to modify any of them?"
         
-        self.context.add_to_history("assistant", chat_response)
-        return chat_response, formatted_objectives
-    
-    def generate_fallback_objectives(self, user_input: str) -> str:
-        """Generate basic objectives if the main generation fails"""
-        return f"""# M1
-Add Main mission_objective
-- Complete the primary mission goal
-Add Secondary bonus_objective
-- Optional additional task
-#"""
+        return chat_response, output
 
 def create_gradio_interface():
     generator = MissionGenerator()
@@ -101,13 +98,14 @@ def create_gradio_interface():
     with gr.Blocks() as interface:
         gr.Markdown("""
         # Original War Mission Objective Generator
-        Describe your mission scenario in natural language, and I'll help you create formatted mission objectives.
+        Describe your mission scenario and I'll help you create formatted mission objectives.
+        Try phrases like "The player must...", "They can optionally...", or "Alternatively..."
         """)
         
         chatbot = gr.Chatbot(height=400, type="messages")
         msg = gr.Textbox(
             label="Describe your mission scenario",
-            placeholder="Tell me about the mission you want to create..."
+            placeholder="Describe what the player needs to do..."
         )
         clear = gr.Button("Clear Conversation")
         formatted_output = gr.Textbox(
@@ -123,9 +121,9 @@ def create_gradio_interface():
         
         gr.Examples(
             examples=[
-                "I need a mission where players have to infiltrate an enemy base. They should try to avoid detection, but if they get spotted, they'll need to fight their way through.",
-                "Create a defensive mission where players protect a convoy. They should also try to minimize civilian casualties.",
-                "I want players to capture a strategic point. They can either do it by force or try diplomatic negotiations with the local faction."
+                "The player must infiltrate the enemy base. They can optionally rescue prisoners along the way. Alternatively, they could create a diversion at the front gate.",
+                "Protect the convoy until it reaches the destination. You can also secure supply drops if possible.",
+                "Your main goal is to capture the strategic point. Alternatively, try negotiating with the local faction silently.",
             ],
             inputs=msg
         )
@@ -135,35 +133,4 @@ def create_gradio_interface():
 # Launch the interface
 if __name__ == "__main__":
     iface = create_gradio_interface()
-    iface.launch()
-
-"""
-# Discord bot implementation using the same generator
-import discord
-from discord.ext import commands
-import os
-
-class MissionBot(commands.Bot):
-    def __init__(self):
-        super().__init__(command_prefix="!")
-        self.generator = MissionGenerator()
-
-    async def on_ready(self):
-        print(f'{self.user} has connected to Discord!')
-
-    @commands.command(name='mission')
-    async def generate_mission(self, ctx, *, description: str):
-        chat_response, formatted_output = self.generator.generate_response(description)
-        
-        # Split response if it's too long for Discord
-        if len(formatted_output) > 1990:  # Discord has 2000 char limit
-            await ctx.send(f"💭 {chat_response}")
-            await ctx.send(f"```\n{formatted_output}\n```")
-        else:
-            await ctx.send(f"💭 {chat_response}\n\n```\n{formatted_output}\n```")
-
-# Initialize and run the bot
-if __name__ == "__main__":
-    bot = MissionBot()
-    bot.run(os.getenv('DISCORD_TOKEN'))
-"""
+    iface.launch()