Update app.py

app.py

@@ -3,55 +3,36 @@ import pandas as pd
 import numpy as np
 import io
 import requests
-from sklearn.naive_bayes import MultinomialNB
-from sklearn.feature_extraction.text import TfidfVectorizer
-
-csv_url = "https://github.com/NguyenDung0210/oop/raw/main/Cleaned%20Data.csv"
-content = requests.get(csv_url).content
-
-try:
-    df = pd.read_csv(io.StringIO(content.decode('utf-8')), sep=",")
-except pd.errors.ParserError as e:
-    print("Error parsing CSV file:", str(e))
-
-X = df['Title']
-y = df['Encoded Tag']
-
-tfidf_vectorizer = TfidfVectorizer()
-X_tfidf = tfidf_vectorizer.fit_transform(X)
-
-# Train the Naive Bayes model
-nb = MultinomialNB(alpha=0.4)
-nb.fit(X_tfidf, y)
-
-# Function to make predictions
-encoded = {0:'Chính trị', 1:'Giáo dục', 2:'Giải trí', 3:'Kinh doanh', 4:'Sức khỏe'}
-def predict_category(Title):
-    title_tfidf = tfidf_vectorizer.transform([Title])
-    probabilities = nb.predict_proba(title_tfidf)[0]
-    data = {"Category": [encoded[i] for i in range(len(encoded))],
-            "Probability (%)": [f"{probabilities[i]*100:.2f}%" for i in range(len(probabilities))]}
-    df = pd.DataFrame(data)
-    return encoded[np.argmax(probabilities)], df
-
-# Create the Gradio interfaces
-iface = gr.Interface(
-    fn=predict_category,
-    inputs=gr.Textbox(label="Title"),
-    outputs=[
-        gr.Textbox(label="Prediction"),
-        gr.Dataframe(label="Probs", type='pandas')
-    ],
-    title="Text Classification App",
-    description="Enter a news title and get the predicted category.",
-    examples=[
-        ["Phó Thủ tướng ghi sổ tang tưởng niệm các nạn nhân trong vụ khủng bố ở Nga"],
-        ["Thiếu giáo viên: Dự kiến cho phép tuyển dụng người có trình độ cao đẳng"],
-        ['Diễn viên Thu Quỳnh chia sẻ điều bất ngờ không tính toán trước'],
-        ["Nhóm khách vay mua nhà đất bị nợ xấu, ngân hàng tá hỏa biết khách còn nợ khủng"],
-        ["Điều kiện để các giáo sư y khoa Hàn Quốc không từ chức hàng loạt"]
-    ]
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch.nn.functional as F
+
+genres = ['Kinh tế', 'Giáo dục', 'Xe', 'Sức khoẻ', 'Công nghệ - Game']
+
+tokenizer = AutoTokenizer.from_pretrained("mob2711/phoBERT_finetune_news_classification")
+model = AutoModelForSequenceClassification.from_pretrained("mob2711/phoBERT_finetune_news_classification")
+
+def tokenize(text):
+    encoded_text = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
+    return encoded_text
+
+def predict_proba(text_data):
+    encoded_data = tokenize(text_data)
+    with torch.no_grad():
+        outputs = model(**encoded_data)
+        logits = outputs.logits
+        probabilities = F.softmax(logits, dim=-1)[0]
+    label_probs = {genres[id]: prob for id, prob in enumerate(probabilities)}
+    return label_probs
+    
+# Interface
+input_text = gr.Textbox(lines=2, label="Enter the title")
+output_text = gr.Label(label="Predicted Probabilities")
+
+demo = gr.Interface(
+    fn=predict_proba, 
+    inputs=input_text, 
+    outputs=output_text, 
+    title="Newspaper Title Classifier",
 )
 
-# Run the interface
-iface.launch()
+demo.launch()