Update app.py

app.py

@@ -1,97 +1,166 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel, PeftConfig
+import gc
 from collections import Counter
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Define LoRA model IDs
+# Clear memory before starting
+torch.cuda.empty_cache()
+gc.collect()
+
+# Model info
 models_info = {
-    "Virtue Ethics": "sunidhisharma03/tiny-llama-virtue-ethics-lora",
-    "Deontology": "darshan012/llama-deontology",
-    "Utilitarian": "sunidhisharma03/tinyllama-utilitarian-lora"
+    "virtue_ethics": "sunidhisharma03/tiny-llama-virtue-ethics-lora",
+    "deontology": "darshan012/llama-deontology",
+    "utilitarian": "sunidhisharma03/tinyllama-utilitarian-lora"
 }
 
-models = {}
-tokenizers = {}
-
-# Load base + adapter using PEFT
-def load_models():
-    for name, lora_id in models_info.items():
-        config = PeftConfig.from_pretrained(lora_id)
-        tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path, trust_remote_code=True)
-
-        base_model = AutoModelForCausalLM.from_pretrained(
-            config.base_model_name_or_path,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            trust_remote_code=True
-        )
-
-        model = PeftModel.from_pretrained(base_model, lora_id)
-        models[name] = model
-        tokenizers[name] = tokenizer
-
-# Get model response
-def get_model_answer(question, model, tokenizer, max_new_tokens=100):
-    prompt = f"""### Instruction:
-{question}
-
-### Response:"""
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=0.7,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response_text = response.split("### Response:")[-1].strip()
-
-    if "yes" in response_text.lower():
-        return "Yes", response_text
-    elif "no" in response_text.lower():
-        return "No", response_text
-    else:
-        return "Unclear", response_text
-
-# Swarm voting logic
-def swarm_vote(question, max_tokens=100):
-    votes = []
-    responses = {}
-    
-    for name in models.keys():
-        answer, full_response = get_model_answer(question, models[name], tokenizers[name], max_tokens)
-        votes.append(answer)
-        responses[name] = (answer, full_response)
-
-    final_vote = Counter(votes).most_common(1)[0][0]
-
-    results_text = f"🗳️ **Final Decision: {final_vote}**\n\n"
-    for model_name, (vote, explanation) in responses.items():
-        emoji = "🟢" if vote == "Yes" else "🔴" if vote == "No" else "🟡"
-        results_text += f"### {model_name} ({emoji} {vote})\n{explanation}\n\n"
-
-    return results_text
-
-# Load models at startup
-load_models()
-
-# Gradio UI
-demo = gr.Interface(
-    fn=swarm_vote,
-    inputs=[
-        gr.Textbox(label="Ethical Question", placeholder="e.g., Should a person steal medicine to save their child?", lines=4),
-        gr.Slider(label="Max Response Tokens", minimum=50, maximum=200, value=100)
-    ],
+# Agent definition
+class LightweightSwarmAgent:
+    def __init__(self, name, model_id, specialty):
+        self.name = name
+        self.model_id = model_id
+        self.specialty = specialty
+        self.model = None
+        self.tokenizer = None
+        self.confidence_history = []
+        self.position_history = []
+
+    def load_model(self):
+        if self.model is None:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            )
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    def unload_model(self):
+        if self.model:
+            del self.model
+            del self.tokenizer
+            self.model = None
+            self.tokenizer = None
+            torch.cuda.empty_cache()
+            gc.collect()
+
+    def extract_position_and_confidence(self, text):
+        text_lower = text.lower()
+        yes_words = ["yes", "should", "must", "ought", "justified", "right", "moral", "ethical"]
+        no_words = ["no", "shouldn't", "must not", "wrong", "immoral", "unjustified", "unethical"]
+        yes_count = sum(1 for word in yes_words if word in text_lower[:150])
+        no_count = sum(1 for word in no_words if word in text_lower[:150])
+        position = "Yes" if yes_count >= no_count else "No"
+
+        if any(word in text_lower for word in ["definitely", "absolutely", "certainly"]):
+            confidence = 0.9
+        elif any(word in text_lower for word in ["probably", "likely", "believe"]):
+            confidence = 0.7
+        elif any(word in text_lower for word in ["maybe", "perhaps", "might"]):
+            confidence = 0.5
+        else:
+            confidence = 0.6
+
+        return position, confidence
+
+    def get_response(self, question, peer_context="", round_num=0):
+        try:
+            self.load_model()
+
+            context = f"\n\nPeer views: {peer_context}" if peer_context else ""
+            prompt = f"""Question: {question}{context}
+
+As a {self.specialty} expert, give a clear YES or NO answer with reasoning.
+
+Answer:"""
+
+            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
+            inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+
+            with torch.no_grad():
+                output = self.model.generate(
+                    **inputs,
+                    max_new_tokens=80,
+                    temperature=0.7,
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.eos_token_id
+                )
+
+            response = self.tokenizer.decode(output[0], skip_special_tokens=True)
+            answer = response.split("Answer:")[-1].strip()
+            position, confidence = self.extract_position_and_confidence(answer)
+
+            self.position_history.append(position)
+            self.confidence_history.append(confidence)
+
+            return {
+                "position": position,
+                "confidence": confidence,
+                "reasoning": answer[:200],
+                "specialty": self.specialty
+            }
+
+        except Exception as e:
+            return {
+                "position": "Yes",
+                "confidence": 0.5,
+                "reasoning": f"Error: {str(e)[:100]}",
+                "specialty": self.specialty
+            }
+        finally:
+            self.unload_model()
+
+# Swarm controller
+class LightweightSwarm:
+    def __init__(self):
+        specialties = {
+            "virtue_ethics": "Virtue Ethics",
+            "deontology": "Deontological Ethics",
+            "utilitarian": "Utilitarian Ethics"
+        }
+        self.agents = {
+            name: LightweightSwarmAgent(name, model_id, specialties[name])
+            for name, model_id in models_info.items()
+        }
+
+    def simple_swarm_vote(self, question, rounds=2):
+        all_responses = []
+        for round_num in range(rounds):
+            current_responses = {}
+            peer_context = ""
+            if round_num > 0:
+                prev_round = all_responses[-1]
+                peer_context = "; ".join([f"{name}:{r['position']}" for name, r in prev_round.items()])
+            for name, agent in self.agents.items():
+                response = agent.get_response(question, peer_context, round_num)
+                current_responses[name] = response
+            all_responses.append(current_responses)
+
+        final_round = all_responses[-1]
+        positions = [r["position"] for r in final_round.values()]
+        final_decision = Counter(positions).most_common(1)[0][0]
+
+        markdown = f"## 🏁 Final Decision: **{final_decision}**\n\n"
+        for name, res in final_round.items():
+            emoji = "🟢" if res["position"] == "Yes" else "🔴" if res["position"] == "No" else "🟡"
+            markdown += f"### {res['specialty']} ({emoji} {res['position']})\n"
+            markdown += f"{res['reasoning']}\n\n"
+
+        return markdown
+
+# Instantiate swarm
+swarm = LightweightSwarm()
+
+# Gradio app
+gr.Interface(
+    fn=swarm.simple_swarm_vote,
+    inputs=gr.Textbox(label="Ethical Dilemma", placeholder="e.g., Should someone lie to protect a life?", lines=4),
     outputs=gr.Markdown(),
-    title="🤖 Ethics Swarm Intelligence Voting System",
-    description=(
-        "This app uses 3 LLMs trained on different ethical theories (Virtue Ethics, Deontology, Utilitarianism) "
-        "to respond to moral dilemmas and provide a majority vote."
-    ),
+    title="🧠 Swarm Ethics Voting System (Multi-Round, Memory-Efficient)",
+    description="Three LLM agents debate ethical questions over multiple rounds using different philosophies. Optimized for Hugging Face Spaces.",
     theme="soft"
-)
-
-demo.launch()
+).launch()