Spaces:
Sleeping
Sleeping
from typing import List | |
from transformers import ( | |
pipeline, | |
AutoTokenizer, | |
AutoModelForSequenceClassification, | |
AutoModelForTokenClassification, | |
) | |
from bs4 import BeautifulSoup | |
import requests | |
# --------------------------------------------------------------------------- | |
# Model identifiers – use your custom sentiment model hosted on Hugging Face | |
# --------------------------------------------------------------------------- | |
SENTIMENT_MODEL_ID = "LinkLinkWu/Stock_Analysis_Test_Ahamed" # binary sentiment | |
NER_MODEL_ID = "dslim/bert-base-NER" | |
# --------------------------------------------------------------------------- | |
# Eager initialisation (singletons shared by the whole Streamlit session) | |
# --------------------------------------------------------------------------- | |
# Sentiment pipeline – returns one label with its score. We will *ignore* the | |
# numeric score down‑stream to satisfy the "no numbers" requirement. | |
sentiment_tokenizer = AutoTokenizer.from_pretrained(SENTIMENT_MODEL_ID) | |
sentiment_model = AutoModelForSequenceClassification.from_pretrained(SENTIMENT_MODEL_ID) | |
sentiment_pipeline = pipeline( | |
"sentiment-analysis", | |
model=sentiment_model, | |
tokenizer=sentiment_tokenizer, | |
) | |
# Named‑entity‑recognition pipeline (ORG extraction) | |
ner_tokenizer = AutoTokenizer.from_pretrained(NER_MODEL_ID) | |
ner_model = AutoModelForTokenClassification.from_pretrained(NER_MODEL_ID) | |
ner_pipeline = pipeline( | |
"ner", | |
model=ner_model, | |
tokenizer=ner_tokenizer, | |
grouped_entities=True, | |
) | |
# --------------------------------------------------------------------------- | |
# Web‑scraping helper (Finviz) | |
# --------------------------------------------------------------------------- | |
def fetch_news(ticker: str) -> List[dict]: | |
"""Return at most 30 latest Finviz headlines for *ticker* ("title" & "link").""" | |
try: | |
url = f"https://finviz.com/quote.ashx?t={ticker}" | |
headers = { | |
"User-Agent": "Mozilla/5.0", | |
"Accept": "text/html", | |
"Accept-Language": "en-US,en;q=0.5", | |
"Referer": "https://finviz.com/", | |
"Connection": "keep-alive", | |
} | |
r = requests.get(url, headers=headers, timeout=10) | |
if r.status_code != 200: | |
return [] | |
soup = BeautifulSoup(r.text, "html.parser") | |
if ticker.upper() not in (soup.title.text if soup.title else "").upper(): | |
return [] # possibly a redirect page | |
table = soup.find(id="news-table") | |
if table is None: | |
return [] | |
headlines: List[dict] = [] | |
for row in table.find_all("tr")[:30]: | |
link_tag = row.find("a") | |
if link_tag: | |
headlines.append({"title": link_tag.get_text(strip=True), "link": link_tag["href"]}) | |
return headlines | |
except Exception: | |
return [] | |
# --------------------------------------------------------------------------- | |
# Sentiment helpers – binary classification, *no* numeric score exposed | |
# --------------------------------------------------------------------------- | |
_LABEL_MAP = {"LABEL_0": "Negative", "LABEL_1": "Positive"} # adjust if model config differs | |
def analyze_sentiment(text: str, pipe=None) -> str: | |
"""Return **"Positive"** or **"Negative"** for a single headline. | |
*Neutral* outputs (if ever returned by the model) are coerced to *Negative*. | |
Numeric confidence scores are deliberately discarded to honour the | |
"no numbers" requirement. | |
""" | |
try: | |
sentiment_pipe = pipe or sentiment_pipeline | |
result = sentiment_pipe(text, truncation=True, return_all_scores=False)[0] | |
raw_label = result.get("label", "").upper() | |
label = _LABEL_MAP.get(raw_label, "Negative") # default to Negative | |
return label | |
except Exception: | |
return "Unknown" | |
# --------------------------------------------------------------------------- | |
# Aggregation – majority vote (Positive‑ratio) → binary label | |
# --------------------------------------------------------------------------- | |
_POS_RATIO_THRESHOLD = 0.6 # ≥60 % positives → overall Positive | |
def aggregate_sentiments(labels: List[str], pos_ratio_threshold: float = _POS_RATIO_THRESHOLD) -> str: | |
"""Combine individual headline labels into an overall binary sentiment. | |
* If *Positive* proportion ≥ *pos_ratio_threshold* → *Positive*. | |
* Otherwise → *Negative*. | |
* Empty list → *Unknown*. | |
""" | |
if not labels: | |
return "Unknown" | |
total = len(labels) | |
positives = sum(1 for l in labels if l == "Positive") | |
ratio = positives / total | |
return "Positive" if ratio >= pos_ratio_threshold else "Negative" | |
# --------------------------------------------------------------------------- | |
# ORG‑entity extraction (ticker discovery) | |
# --------------------------------------------------------------------------- | |
def extract_org_entities(text: str, pipe=None, max_entities: int = 5) -> List[str]: | |
"""Extract up to *max_entities* unique ORG tokens (upper‑case, de‑hashed).""" | |
try: | |
ner_pipe = pipe or ner_pipeline | |
entities = ner_pipe(text) | |
orgs: List[str] = [] | |
for ent in entities: | |
if ent.get("entity_group") == "ORG": | |
token = ent["word"].replace("##", "").strip().upper() | |
if token and token not in orgs: | |
orgs.append(token) | |
if len(orgs) >= max_entities: | |
break | |
return orgs | |
except Exception: | |
return [] | |
# --------------------------------------------------------------------------- | |
# Public accessors (legacy compatibility) | |
# --------------------------------------------------------------------------- | |
def get_sentiment_pipeline(): | |
return sentiment_pipeline | |
def get_ner_pipeline(): | |
return ner_pipeline | |