Setup the app code, requirements, and metadata

README.md

@@ -1,10 +1,10 @@
 ---
-title: Streaming Llm Weather Alerts
-emoji: ⚡
-colorFrom: green
-colorTo: red
+title: Streaming LLM Weather Alerts
+emoji: 🌤️
+colorFrom: pink
+colorTo: gray
 sdk: gradio
-sdk_version: 4.40.0
+sdk_version: 4.36.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -0,0 +1,139 @@
+import apache_beam as beam
+import gradio as gr
+import huggingface_hub
+import pandas as pd
+import plotly.graph_objects as go
+import spaces
+import textwrap
+import torch
+import us
+
+from apache_beam.options.pipeline_options import PipelineOptions, SetupOptions
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+import json
+import logging
+import os
+import requests
+
+MODEL_NAME = "google/gemma-2-2b-it"
+PROMPT_TEMPLATE = """Write a succinct summary of the following weather alerts. Do not comment on missing information - just summarize the information provided/available.
+
+```json
+{}
+```
+
+Summary (In the state...):
+"""
+
+# Initialize an empty list to store weather alerts
+alerts = []
+
+
+# Define a transform for fetching weather alerts
+class FetchWeatherAlerts(beam.DoFn):
+    def process(self, state):
+        logging.info(f"Fetching weather alerts for {state} from weather.gov")
+        url = f"https://api.weather.gov/alerts/active?area={state}"
+        response = requests.get(
+            url,
+            headers={
+                "User-Agent": "(Neal DeBuhr, https://huggingface.co/spaces/ndebuhr/streaming-llm-weather-alerts)",
+                "Accept": "application/geo+json",
+            },
+        )
+        if response.status_code == 200:
+            logging.info(f"Fetched weather alerts for {state} from weather.gov")
+            features = response.json()["features"]
+            alerts.append(
+                {
+                    "features": [
+                        {
+                            "event": feature["properties"]["event"],
+                            "headline": feature["properties"]["headline"],
+                            "instruction": feature["properties"]["instruction"],
+                        }
+                        for feature in features
+                        if feature["properties"]["messageType"] == "Alert"
+                    ],
+                    "state": state,
+                }
+            )
+
+
+pipeline_options = PipelineOptions()
+# Save the main session state so that pickled functions and classes
+# defined in __main__ can be unpickled
+pipeline_options.view_as(SetupOptions).save_main_session = True
+
+# Create and run the Apache Beam pipeline to fetch weather alerts
+with beam.Pipeline(options=pipeline_options) as p:
+    (p
+        | "Create States" >> beam.Create([state.abbr for state in us.states.STATES])
+        | "Fetch Weather Alerts" >> beam.ParDo(FetchWeatherAlerts())
+    )
+
+
+# Define a function to generate alert summaries using transformers and ZeroGPU
+@spaces.GPU(duration=300)
+def generate_summaries(alerts):
+    huggingface_hub.login(token=os.environ["HUGGINGFACE_TOKEN"])
+    device = torch.device("cuda")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
+    for alert in alerts:
+        prompt = PROMPT_TEMPLATE.format(json.dumps(alert, indent=2))
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs, max_new_tokens=256, pad_token_id=tokenizer.eos_token_id
+            )
+
+        alert["summary"] = (
+            tokenizer.decode(outputs[0], skip_special_tokens=True)
+            .replace(prompt, "")
+            .strip()
+        )
+    return alerts
+
+
+alerts = generate_summaries(alerts)
+
+df = pd.DataFrame.from_dict(
+    [{"state": alert["state"], "summary": alert["summary"]} for alert in alerts]
+)
+
+
+def get_map():
+    def wrap_text(text, width=50):
+        return "<br>".join(textwrap.wrap(text, width=width))
+
+    df["wrapped_summary"] = df["summary"].apply(wrap_text)
+
+    fig = go.Figure(
+        go.Choropleth(
+            locations=df["state"],
+            z=[1 for _ in df["summary"]],
+            locationmode="USA-states",
+            colorscale=[
+                [0, "lightgrey"],
+                [1, "lightgrey"],
+            ],  # Single color for all states
+            showscale=False,
+            text=df["wrapped_summary"],
+            hoverinfo="text",
+            hovertemplate="%{text}<extra></extra>",
+        )
+    )
+
+    fig.update_layout(title_text="Streaming LLM Weather Alerts", geo_scope="usa")
+
+    return fig
+
+
+# Create Gradio interface
+iface = gr.Interface(fn=get_map, inputs=None, outputs=gr.Plot())
+
+# Launch the Gradio interface
+iface.launch()

requirements.txt

@@ -0,0 +1,6 @@
+apache_beam==2.57.0
+huggingface_hub==0.24.5
+pandas==2.2.2
+plotly==5.23.0
+transformers==4.43.4
+us==3.2.0