Update tools/news_reporter.py

tools/news_reporter.py

@@ -1,39 +1,42 @@
+
 from agent_state import AgentState
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import os
-from langchain_huggingface import HuggingFaceEndpoint
-
-HF_TOKEN = os.getenv("HF_TOKEN")
-if not HF_TOKEN:
-    raise RuntimeError("HF_TOKEN env var is missing")
 
 MODEL_ID = "google/gemma-3n-E4B-it"   # confirm this exact repo id exists
 
-llm = HuggingFaceEndpoint(
-    repo_id=MODEL_ID,
-    task="text-generation",
-    huggingfacehub_api_token=HF_TOKEN,
-    temperature=0.7,
+# Optional: use a cache directory
+cache_dir = os.getenv("HF_HOME", None)
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID, cache_dir=cache_dir
+)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, cache_dir=cache_dir
+)
+
+# Create text generation pipeline
+text_generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
     max_new_tokens=512,
-    # or: model_kwargs={"temperature":0.7, "max_new_tokens":512}
+    temperature=0.7,
 )
 
 def create_news_report(state: AgentState) -> AgentState:
     if state.feedback:
-        prompt = (
-            "You are revising a news report based on the user's feedback.\n"
-            f'Transcription: "{state.transcribed_text}"\n'
-            f'Old Report: "{state.news_report}"\n'
-            f'Feedback: "{state.feedback}"\n'
-            "Rewrite it professionally."
-        )
+        prompt = f"""You are revising a news report based on the user's feedback:
+                    Transcription: "{state.transcribed_text}"
+                    Old Report: "{state.news_report}"
+                    Feedback: "{state.feedback}" """
     else:
-        prompt = (
-            "Write a professional news article based on this transcription:\n"
-            f'"{state.transcribed_text}"'
-        )
+        prompt = f"""Write a professional news article based on this transcription:
+                    "{state.transcribed_text}" """
 
-    report = llm.invoke(prompt).strip()
-    state.news_report = report
+    generated = text_generator(prompt)[0]["generated_text"]
+    state.news_report = generated
     state.feedback = None
     state.approved = False
     return state