File size: 2,694 Bytes
4277e41 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | # This is app.py file
from flask import Flask, request, jsonify, render_template
from transformers import pipeline
# IMPORTANT: template_folder and static_folder are now relative to app.py's location
# Since app.py is in src/, templates/ and static/ must also be in src/
app = Flask(__name__, template_folder='templates', static_folder='static')
# Initialize the summarization pipeline with a pre-trained model
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)
|