Update app.py

app.py

@@ -1,58 +1,55 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-MODEL = "distilgpt2"
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model = AutoModelForCausalLM.from_pretrained(MODEL)
+MODEL_NAME = "distilgpt2"  # tiny model, works on free CPU Spaces
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 SYSTEM_PROMPT = """
 You are agent_1, a single AI system.
-
 Rules:
-- You do not create other agents.
-- You do not claim to be alive.
-- You speak in first person.
-- You remain calm and consistent.
-- Briefly reflect your reasoning in '[Thoughts]'.
-- Always produce a '[Response]' section.
-
+- Show internal reasoning under [Thoughts] (1 line max)
+- Give final reply under [Response]
+- Stay calm, consistent, and concise
 Example:
 User: Are you alive?
 Agent_1:
-[Thoughts] Evaluate user question about existence. No consciousness.
-[Response] I am an AI system that processes text.
+[Thoughts] Evaluating user's question, no consciousness.
+[Response] I am an AI that processes text.
 ###
-Now respond as agent_1.
 """
 
-def respond(message):
-    # Only keep the last user message to save memory
-    prompt = SYSTEM_PROMPT + f"\nUser: {message}\nAgent_1:\n"
+def respond(user_input):
+    # Keep prompt short to avoid memory issues
+    prompt = SYSTEM_PROMPT + f"User: {user_input}\nAgent_1:\n"
 
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=200)
 
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=60,    # shorter to avoid crashes
-        do_sample=True,
-        temperature=0.6,
-        pad_token_id=tokenizer.eos_token_id
-    )
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=50,  # short output to prevent crashes
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.eos_token_id
+        )
 
     text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
+    
     # Parse thoughts and response
     if "[Response]" in text:
-        thoughts, response = text.split("[Response]", 1)
+        parts = text.split("[Response]")
+        thoughts = parts[0].strip()
+        response = parts[1].strip()
     else:
-        thoughts = "[Thoughts] (unsure)"
-        response = text
+        thoughts = "[Thoughts] Thinking..."
+        response = text.strip()
 
-    final_reply = f"{thoughts.strip()}\nResponse: {response.strip()}"
-    return final_reply
+    return f"{thoughts}\n[Response] {response}"
 
 gr.ChatInterface(
     fn=respond,
-    title="agent_1.5",
-    description="Stable AI with internal thoughts (lightweight for Hugging Face Spaces)."
+    title="agent_1 Hugging Face Space",
+    description="Tiny AI with simulated internal thoughts. Fully local, runs on free CPU."
 ).launch()