Update app.py

app.py

@@ -1,46 +1,49 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# 1. Load YOUR Model
+# 1. Load the GOOD model (It's safe on main!)
 MODEL_NAME = "himu1780/ai-python-model"
 
 print(f"Loading {MODEL_NAME}...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-# Create a text generation pipeline
+# 2. Create the pipeline
 generator = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=-1 # Use CPU
+    device=-1 # CPU
 )
 
 def chat_function(user_message, history):
-    # 2. Format the prompt
+    # 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. Generate the answer
+    # 4. Generate with Penalties (Bans looping)
     response = generator(
         prompt,
-        max_new_tokens=200,
-        temperature=0.6,
+        max_new_tokens=250,
+        temperature=0.5,        # Lower = More logical/focused
+        repetition_penalty=1.2, # <--- Kills the "Hello World" loops
         do_sample=True,
         pad_token_id=tokenizer.eos_token_id
     )
     
-    # 4. Clean up the text
+    # 5. Clean output
     full_text = response[0]['generated_text']
+    # Only keep the part after "Response:"
     answer = full_text.split("### Response:")[-1].strip()
     
     return answer
 
-# 5. Launch the Chat UI (Fixed: Removed 'theme' argument)
+# 6. Launch
 demo = gr.ChatInterface(
     fn=chat_function,
-    title="🐍 AI Python Assistant",
-    description="Ask me to write Python code!",
-    examples=["Write a function to calculate area of circle", "Create a list of numbers 1 to 10"]
+    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"]
 )
 
 demo.launch()