| |
| from flask import Flask, request, jsonify, render_template |
| from transformers import pipeline |
| |
| |
| app = Flask(__name__, template_folder='templates', static_folder='static') |
|
|
| |
| try: |
| pipe = pipeline("text2text-generation", model="t5-small") |
| print("AI Model 't5-small' loaded successfully!") |
| except Exception as e: |
| print(f"Error loading AI model: {e}") |
| pipe = None |
|
|
| @app.route('/') |
| def home(): |
| """ |
| Renders the main index.html page. |
| """ |
| return render_template('index.html') |
|
|
| @app.route('/rewrite', methods=['POST']) |
| def rewrite_message(): |
| """ |
| API endpoint to rewrite a message using the AI model. |
| Expects JSON input with 'message' and 'tone'. |
| """ |
| if pipe is None: |
| return jsonify({"error": "AI model not loaded. Please try again later."}), 500 |
|
|
| data = request.get_json() |
| message = data.get('message') |
| tone = data.get('tone', 'neutral') |
|
|
| if not message: |
| return jsonify({"error": "No message provided."}), 400 |
|
|
| prompt = f"Rewrite the following message to be more {tone}: '{message}'" |
|
|
| try: |
| rewritten_output = pipe(prompt, max_length=150, min_length=20, num_beams=5, no_repeat_ngram_size=2) |
| rewritten_text = rewritten_output[0]['generated_text'] |
|
|
| original_tone_analysis = "Original tone perceived as: " + detect_simple_tone(message) |
| suggested_action = "Consider if the new tone perfectly conveys your intent." |
|
|
| return jsonify({ |
| "rewrittenText": rewritten_text, |
| "toneAnalysis": original_tone_analysis, |
| "suggestion": suggested_action |
| }) |
|
|
| except Exception as e: |
| print(f"Error during AI generation: {e}") |
| return jsonify({"error": f"Failed to rewrite message: {str(e)}"}), 500 |
|
|
| def detect_simple_tone(text): |
| """ |
| A very basic function to 'detect' original tone. |
| """ |
| text_lower = text.lower() |
| if any(word in text_lower for word in ["angry", "frustrated", "bad", "hate", "unacceptable"]): |
| return "Negative/Frustrated" |
| elif any(word in text_lower for word in ["happy", "great", "excellent", "good", "positive"]): |
| return "Positive" |
| elif any(word in text_lower for word in ["demand", "must", "require"]): |
| return "Demanding" |
| elif any(word in text_lower for word in ["please", "kindly", "suggest"]): |
| return "Polite/Suggestive" |
| return "Neutral" |
|
|
| if __name__ == '__main__': |
| app.run(host='0.0.0.0', port=7860) |
| |
