Update app.py

app.py

@@ -1,49 +1,62 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# 1. Load the GOOD model (It's safe on main!)
+# 1. Load Your Smart Model (Step 1000)
 MODEL_NAME = "himu1780/ai-python-model"
 
 print(f"Loading {MODEL_NAME}...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-# 2. Create the pipeline
+# 2. Create the Generator Pipeline
 generator = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=-1 # CPU
+    device=-1  # Use CPU
 )
 
 def chat_function(user_message, history):
-    # 3. THE MAGIC FIX: Format the input exactly how the model likes it
-    # We add "def" at the end to force it to start coding immediately
-    prompt = f"### Instruction:\n{user_message}\n\n### Response:\n"
+    # 3. THE FIX: Force the prompt structure
+    # We add "```python" at the end. This TRICKS the model.
+    # It thinks: "Oh, I already started writing a code block, I must finish it."
+    prompt = f"### Instruction:\n{user_message}\n\n### Response:\n```python\n"
     
-    # 4. Generate with Penalties (Bans looping)
+    # 4. Generate with Anti-Loop Settings
     response = generator(
         prompt,
-        max_new_tokens=250,
-        temperature=0.5,        # Lower = More logical/focused
-        repetition_penalty=1.2, # <--- Kills the "Hello World" loops
+        max_new_tokens=250,      # Give it enough space to write a full function
+        temperature=0.4,         # Low temperature = Focused logic (No creative babbling)
+        repetition_penalty=1.2,  # STRICT penalty. If it repeats "Hello", it gets banned.
         do_sample=True,
         pad_token_id=tokenizer.eos_token_id
     )
     
-    # 5. Clean output
+    # 5. Clean up the output
     full_text = response[0]['generated_text']
-    # Only keep the part after "Response:"
-    answer = full_text.split("### Response:")[-1].strip()
     
-    return answer
+    # Isolate just the new code
+    # We split by "### Response:" to get the answer
+    answer_part = full_text.split("### Response:")[-1].strip()
+    
+    # Ensure the formatting looks nice in the chat
+    # We stripped the ```python start, so we add it back for the UI
+    if not answer_part.startswith("```"):
+        answer_part = "```python\n" + answer_part
+        
+    return answer_part
 
-# 6. Launch
+# 6. Launch the Chat Interface
 demo = gr.ChatInterface(
     fn=chat_function,
-    title="🐍 AI Python Graduate (Step 1000)",
-    description="I am fully trained! Ask me to 'Write a function to...' or 'Create a loop...'",
-    examples=["Write a python function to add two numbers", "Create a loop from 1 to 10"]
+    title="🐍 AI Python Assistant",
+    description="I am trained! Ask me to 'Write a function to...' or 'Create a loop...'",
+    examples=[
+        "Write a python function to add two numbers",
+        "Create a loop from 1 to 10",
+        "Write a script to calculate the area of a circle"
+    ],
+    theme="soft" 
 )
 
 demo.launch()