Update app.py

app.py

@@ -1,153 +1,159 @@
-from flask import Flask, render_template, request, jsonify
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TFAutoModelForSeq2SeqLM
-import torch
-import tensorflow as tf
-import numpy as np
-import os
-
-app = Flask(__name__)
-
-MODELS_CONFIG = {
-    "correction": {"path": "vietnamese-error-correction", "framework": "pt"},
-    "diacritics": {"path": "vietnamese-diacritic-restoration-v2", "framework": "tf"} 
-}
-
-loaded_models = {}
-
-print("Đang khởi tạo các models...")
-device_pt = "cuda" if torch.cuda.is_available() else "cpu"
-
-for mode, config in MODELS_CONFIG.items():
-    path = config["path"]
-    fw = config["framework"]
-    try:
-        print(f"Loading model {mode} ({fw}) từ {path}...")
-        tokenizer = AutoTokenizer.from_pretrained(path)
-        
-        if fw == "pt":
-            model = AutoModelForSeq2SeqLM.from_pretrained(path).to(device_pt)
-        else:
-            model = TFAutoModelForSeq2SeqLM.from_pretrained(path)
-            
-        loaded_models[mode] = {
-            "tokenizer": tokenizer,
-            "model": model,
-            "framework": fw
-        }
-        print(f"Model {mode} đã sẵn sàng!")
-    except Exception as e:
-        print(f"Lỗi khi load model {mode}: {e}")
-
-def process_with_confidence(text, mode):
-    if mode not in loaded_models:
-        raise ValueError(f"Model {mode} chưa được load.")
-    
-    m_info = loaded_models[mode]
-    tokenizer = m_info["tokenizer"]
-    model = m_info["model"]
-    fw = m_info["framework"]
-
-    if fw == "pt":
-        inputs = tokenizer(text, return_tensors="pt").to(device_pt)
-    else:
-        inputs = tokenizer(text, return_tensors="tf")
-
-    if fw == "pt":
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=256,
-                return_dict_in_generate=True,
-                output_scores=True
-            )
-        transition_scores = model.compute_transition_scores(
-            outputs.sequences, outputs.scores, normalize_logits=True
-        )
-        transition_scores = transition_scores.cpu().numpy()
-        generated_tokens = outputs.sequences[0].cpu().numpy()
-    else:
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            return_dict_in_generate=True,
-            output_scores=True
-        )
-        transition_scores = model.compute_transition_scores(
-            outputs.sequences, outputs.scores, normalize_logits=True
-        )
-        transition_scores = transition_scores.numpy()
-        generated_tokens = outputs.sequences[0].numpy()
-
-    special_tokens = {tokenizer.bos_token_id, tokenizer.eos_token_id, tokenizer.pad_token_id}
-    
-    start_index = 0
-    while start_index < len(generated_tokens) and generated_tokens[start_index] in special_tokens:
-        start_index += 1
-    
-    end_index = len(generated_tokens)
-    for i in range(start_index, len(generated_tokens)):
-        if generated_tokens[i] in special_tokens:
-            end_index = i
-            break
-            
-    output_ids = generated_tokens[start_index:end_index]
-    full_text = tokenizer.decode(output_ids, skip_special_tokens=True)
-    target_words = full_text.split()
-
-    if not target_words:
-        return full_text, []
-
-    token_to_word_map = []
-    for i, token_id in enumerate(output_ids):
-        if i >= len(transition_scores[0]): break
-        
-        log_prob = transition_scores[0][i]
-        prob = np.exp(log_prob)
-
-        decoded_up_to_here = tokenizer.decode(output_ids[:i+1], skip_special_tokens=True)
-        words_so_far = decoded_up_to_here.split()
-        word_index = len(words_so_far) - 1 if words_so_far else 0
-        
-        token_to_word_map.append({'prob': prob, 'word_index': word_index})
-
-    word_confidences = {}
-    for item in token_to_word_map:
-        idx = item['word_index']
-        if idx not in word_confidences: word_confidences[idx] = []
-        word_confidences[idx].append(item['prob'])
-
-    confidence_list = []
-    for i in range(len(target_words)):
-        if i in word_confidences:
-            probs = word_confidences[i]
-            confidence_list.append(float(np.mean(probs)))
-        else:
-            confidence_list.append(0.0)
-
-    return full_text, confidence_list
-
-@app.route('/')
-def index():
-    return render_template('index.html')
-
-@app.route('/correct', methods=['POST'])
-def correct_text():
-    data = request.get_json()
-    input_text = data.get('text', '')
-    mode = data.get('mode', 'correction') 
-
-    if not input_text.strip():
-        return jsonify({"result": "", "confidences": []})
-
-    try:
-        generated_text, confidences = process_with_confidence(input_text, mode)
-        return jsonify({
-            "result": generated_text,
-            "confidences": confidences
-        })
-    except Exception as e:
-        print(f"Error: {e}")
-        return jsonify({"error": str(e)}), 500
-
-if __name__ == '__main__':
+from flask import Flask, render_template, request, jsonify
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TFAutoModelForSeq2SeqLM
+import torch
+import tensorflow as tf
+import numpy as np
+import os
+
+app = Flask(__name__)
+
+MODELS_CONFIG = {
+    "correction": {
+        "path": "yammdd/vietnamese-error-correction", 
+        "framework": "pt"
+    },
+    "diacritics": {
+        "path": "yammdd/vietnamese-diacritic-restoration-v2", 
+        "framework": "tf"
+    } 
+}
+
+loaded_models = {}
+
+print("Đang khởi tạo các models...")
+device_pt = "cuda" if torch.cuda.is_available() else "cpu"
+
+for mode, config in MODELS_CONFIG.items():
+    path = config["path"]
+    fw = config["framework"]
+    try:
+        print(f"Loading model {mode} ({fw}) từ {path}...")
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        
+        if fw == "pt":
+            model = AutoModelForSeq2SeqLM.from_pretrained(path).to(device_pt)
+        else:
+            model = TFAutoModelForSeq2SeqLM.from_pretrained(path)
+            
+        loaded_models[mode] = {
+            "tokenizer": tokenizer,
+            "model": model,
+            "framework": fw
+        }
+        print(f"Model {mode} đã sẵn sàng!")
+    except Exception as e:
+        print(f"Lỗi khi load model {mode}: {e}")
+
+def process_with_confidence(text, mode):
+    if mode not in loaded_models:
+        raise ValueError(f"Model {mode} chưa được load.")
+    
+    m_info = loaded_models[mode]
+    tokenizer = m_info["tokenizer"]
+    model = m_info["model"]
+    fw = m_info["framework"]
+
+    if fw == "pt":
+        inputs = tokenizer(text, return_tensors="pt").to(device_pt)
+    else:
+        inputs = tokenizer(text, return_tensors="tf")
+
+    if fw == "pt":
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=256,
+                return_dict_in_generate=True,
+                output_scores=True
+            )
+        transition_scores = model.compute_transition_scores(
+            outputs.sequences, outputs.scores, normalize_logits=True
+        )
+        transition_scores = transition_scores.cpu().numpy()
+        generated_tokens = outputs.sequences[0].cpu().numpy()
+    else:
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            return_dict_in_generate=True,
+            output_scores=True
+        )
+        transition_scores = model.compute_transition_scores(
+            outputs.sequences, outputs.scores, normalize_logits=True
+        )
+        transition_scores = transition_scores.numpy()
+        generated_tokens = outputs.sequences[0].numpy()
+
+    special_tokens = {tokenizer.bos_token_id, tokenizer.eos_token_id, tokenizer.pad_token_id}
+    
+    start_index = 0
+    while start_index < len(generated_tokens) and generated_tokens[start_index] in special_tokens:
+        start_index += 1
+    
+    end_index = len(generated_tokens)
+    for i in range(start_index, len(generated_tokens)):
+        if generated_tokens[i] in special_tokens:
+            end_index = i
+            break
+            
+    output_ids = generated_tokens[start_index:end_index]
+    full_text = tokenizer.decode(output_ids, skip_special_tokens=True)
+    target_words = full_text.split()
+
+    if not target_words:
+        return full_text, []
+
+    token_to_word_map = []
+    for i, token_id in enumerate(output_ids):
+        if i >= len(transition_scores[0]): break
+        
+        log_prob = transition_scores[0][i]
+        prob = np.exp(log_prob)
+
+        decoded_up_to_here = tokenizer.decode(output_ids[:i+1], skip_special_tokens=True)
+        words_so_far = decoded_up_to_here.split()
+        word_index = len(words_so_far) - 1 if words_so_far else 0
+        
+        token_to_word_map.append({'prob': prob, 'word_index': word_index})
+
+    word_confidences = {}
+    for item in token_to_word_map:
+        idx = item['word_index']
+        if idx not in word_confidences: word_confidences[idx] = []
+        word_confidences[idx].append(item['prob'])
+
+    confidence_list = []
+    for i in range(len(target_words)):
+        if i in word_confidences:
+            probs = word_confidences[i]
+            confidence_list.append(float(np.mean(probs)))
+        else:
+            confidence_list.append(0.0)
+
+    return full_text, confidence_list
+
+@app.route('/')
+def index():
+    return render_template('index.html')
+
+@app.route('/correct', methods=['POST'])
+def correct_text():
+    data = request.get_json()
+    input_text = data.get('text', '')
+    mode = data.get('mode', 'correction') 
+
+    if not input_text.strip():
+        return jsonify({"result": "", "confidences": []})
+
+    try:
+        generated_text, confidences = process_with_confidence(input_text, mode)
+        return jsonify({
+            "result": generated_text,
+            "confidences": confidences
+        })
+    except Exception as e:
+        print(f"Error: {e}")
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860, debug=False)