First Model Version

README.md

@@ -1,8 +1,8 @@
 ---
 title: News Analyzer
-emoji: ⚡
-colorFrom: red
-colorTo: blue
+emoji: 📰
+colorFrom: grey
+colorTo: white
 sdk: streamlit
 app_file: app.py
 pinned: false

app.py

@@ -0,0 +1,40 @@
+import streamlit as st
+
+from news_pipeline import NewsPipeline
+
+CATEGORY_EMOJIS = {
+    "Automobile": "🚗",
+    "Entertainment": "🍿",
+    "Politics": "⚖️",
+    "Science": "🧪",
+    "Sports": "🏀",
+    "Technology": "💻",
+    "World": "🌍",
+}
+FAKE_EMOJIS = {"Fake": "👻", "Real": "👍"}
+CLICKBAIT_EMOJIS = {"Clickbait": "🎣", "Normal": "✅"}
+
+
+def app():
+    news_pipe = NewsPipeline()
+    st.title("📰 News Analyzer")
+    headline = st.text_input("Article headline:")
+    content = st.text_area("Article content:")
+    button = st.button("Analyze")
+    if button:
+        with st.spinner("Analyzing article..."):
+            prediction = news_pipe(headline, content)
+        st.success("Article successfully analyzed!")
+        st.markdown(
+            f"{CATEGORY_EMOJIS[prediction['category']]} **Category**: {prediction['category']}"
+        )
+        st.markdown(
+            f"{FAKE_EMOJIS[prediction['fake']]} **Fake**: {'Yes' if prediction['fake'] == 'Fake' else 'No'}"
+        )
+        st.markdown(
+            f"{CLICKBAIT_EMOJIS[prediction['clickbait']]} **Clickbait**: {'Yes' if prediction['clickbait'] == 'Clickbait' else 'No'}"
+        )
+
+
+if __name__ == "__main__":
+    app()

news_pipeline.py

@@ -0,0 +1,44 @@
+from typing import Dict
+
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    TextClassificationPipeline,
+)
+
+
+class NewsPipeline:
+    def __init__(self) -> None:
+        self.category_tokenizer = AutoTokenizer.from_pretrained("elozano/news-category")
+        self.category_pipeline = TextClassificationPipeline(
+            model=AutoModelForSequenceClassification.from_pretrained(
+                "elozano/news-category"
+            ),
+            tokenizer=self.category_tokenizer,
+        )
+        self.fake_tokenizer = AutoTokenizer.from_pretrained("elozano/news-fake")
+        self.fake_pipeline = TextClassificationPipeline(
+            model=AutoModelForSequenceClassification.from_pretrained(
+                "elozano/news-fake"
+            ),
+            tokenizer=self.fake_tokenizer,
+        )
+        self.clickbait_pipeline = TextClassificationPipeline(
+            model=AutoModelForSequenceClassification.from_pretrained(
+                "elozano/news-clickbait"
+            ),
+            tokenizer=AutoTokenizer.from_pretrained("elozano/news-clickbait"),
+        )
+
+    def __call__(self, headline: str, content: str) -> Dict[str, str]:
+        category_article_text = f" {self.category_tokenizer.sep_token} ".join(
+            [headline, content]
+        )
+        fake_article_text = f" {self.fake_tokenizer.sep_token} ".join(
+            [headline, content]
+        )
+        return {
+            "category": self.category_pipeline(category_article_text)[0]["label"],
+            "fake": self.fake_pipeline(fake_article_text)[0]["label"],
+            "clickbait": self.clickbait_pipeline(headline)[0]["label"],
+        }