Update app.py

app.py

@@ -1,63 +1,77 @@
-import os
-from flask import Flask, request, jsonify
-from transformers import RobertaTokenizer, T5ForConditionalGeneration
-import torch
-
-app = Flask(__name__)
-
-# Hugging Face stores files in the current directory
-MODEL_PATH = "./my_codet5_model" 
-
-# Check if model exists
-if not os.path.exists(MODEL_PATH):
-    print(f"❌ Error: Model found at {MODEL_PATH}")
-else:
-    print(f"✅ Model found at {MODEL_PATH}")
-
-device = "cpu" # Free tier is CPU only (but fast enough for this)
-
-print("Loading model...")
-tokenizer = RobertaTokenizer.from_pretrained(MODEL_PATH)
-model = T5ForConditionalGeneration.from_pretrained(MODEL_PATH).to(device)
-print("Model loaded!")
-
-def generate_summary(diff_text):
-    if not diff_text or len(diff_text.strip()) < 5:
-        return "Minor changes"
-    
-    input_text = "Summarize: " + diff_text
-    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
-
-    outputs = model.generate(
-        input_ids, 
-        max_length=50, 
-        num_beams=5, 
-        early_stopping=True
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-@app.route('/generate', methods=['POST'])
-def generate_commit():
-    data = request.json
-    files = data.get('files', [])
-    
-    if not files:
-        return jsonify({"commit_message": ""})
-
-    final_message_parts = []
-    
-    for file_obj in files:
-        name = file_obj.get('name', 'Unknown File')
-        diff = file_obj.get('diff', '')
-        try:
-            summary = generate_summary(diff)
-            final_message_parts.append(f"{name}\n{summary}")
-        except Exception as e:
-            print(f"Error processing {name}: {e}")
-            final_message_parts.append(f"{name}\nUpdate file")
-
-    return jsonify({"commit_message": "\n\n".join(final_message_parts)})
-
-# Hugging Face expects running on port 7860
-if __name__ == '__main__':
+import os
+from flask import Flask, request, jsonify
+from transformers import RobertaTokenizer, T5ForConditionalGeneration
+import torch
+
+app = Flask(__name__)
+
+# --- CRITICAL CHANGE ---
+# Instead of your local folder, we point to a PUBLIC Expert Model
+# This model has read millions of commit messages and knows exactly what to do.
+MODEL_NAME = "ncoop57/commit-t5"
+
+print(f"--- AI Commit Generator Server ---")
+print(f"Downloading/Loading Expert Model: {MODEL_NAME}")
+
+device = "cpu" # HF Spaces free tier is CPU
+
+try:
+    # This will download the model automatically the first time it runs
+    tokenizer = RobertaTokenizer.from_pretrained(MODEL_NAME)
+    model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
+    print("✅ Expert Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    exit(1)
+
+def generate_summary(diff_text):
+    if not diff_text or len(diff_text.strip()) < 5:
+        return "Update file"
+
+    # Preprocess: "commit-t5" just expects the raw code diff, no "Summarize:" prefix needed usually,
+    # but let's keep it simple.
+    input_text = diff_text + " </s>" 
+    
+    input_ids = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True).input_ids.to(device)
+
+    outputs = model.generate(
+        input_ids, 
+        max_length=80,      # Commits are usually short
+        min_length=5,
+        num_beams=5, 
+        early_stopping=True,
+        no_repeat_ngram_size=2 # Stops it from saying "update update update"
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+@app.route('/generate', methods=['POST'])
+def generate_commit():
+    data = request.json
+    files = data.get('files', [])
+    
+    if not files:
+        return jsonify({"commit_message": ""})
+
+    final_message_parts = []
+    
+    for file_obj in files:
+        name = file_obj.get('name', 'Unknown File')
+        diff = file_obj.get('diff', '')
+        
+        # Skip binary files or huge diffs
+        if len(diff) > 4000:
+            final_message_parts.append(f"{name}\nLarge changes detected")
+            continue
+
+        try:
+            summary = generate_summary(diff)
+            # Format: File Name -> The generated message
+            final_message_parts.append(f"{name}\n{summary}")
+        except Exception as e:
+            print(f"Error processing {name}: {e}")
+            final_message_parts.append(f"{name}\nUpdate file")
+
+    return jsonify({"commit_message": "\n\n".join(final_message_parts)})
+
+if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)