Update FinGPT_TaskII_Submission/scripts/inference_with_tools.py

FinGPT_TaskII_Submission/scripts/inference_with_tools.py

@@ -0,0 +1,108 @@
+# ================================================================
+# File: inference_with_tools.py
+# Author: fengxb30
+# Description: Inference script for FinGPT Task II (Compliance Agents)
+# ================================================================
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import requests
+import yfinance as yf
+from datetime import datetime
+import json
+import os
+
+# -----------------------------
+# 1️⃣  Load the model
+# -----------------------------
+def load_model(model_name_or_path="fengxb30/FinGPT_TaskII_Compliance"):
+    print(f"🔹 Loading model from {model_name_or_path} ...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name_or_path,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto"
+    )
+    model.eval()
+    return tokenizer, model
+
+# -----------------------------
+# 2️⃣  Financial Data API Helper
+# -----------------------------
+def fetch_financial_data(ticker="AAPL"):
+    """Fetch real-time market data from Yahoo Finance"""
+    stock = yf.Ticker(ticker)
+    info = stock.info
+    return {
+        "symbol": ticker,
+        "price": info.get("currentPrice", None),
+        "marketCap": info.get("marketCap", None),
+        "fiftyTwoWeekHigh": info.get("fiftyTwoWeekHigh", None),
+        "fiftyTwoWeekLow": info.get("fiftyTwoWeekLow", None),
+    }
+
+# -----------------------------
+# 3️⃣  External Tool: RAG-style retrieval
+# -----------------------------
+def retrieve_context_from_web(query):
+    """Use simple web API (DuckDuckGo Instant Answer) to retrieve context"""
+    url = f"https://api.duckduckgo.com/?q={query}&format=json"
+    try:
+        res = requests.get(url, timeout=10)
+        data = res.json()
+        return data.get("AbstractText") or data.get("Heading") or "No context found."
+    except Exception as e:
+        return f"Retrieval failed: {e}"
+
+# -----------------------------
+# 4️⃣  Model inference with context
+# -----------------------------
+def generate_response(model, tokenizer, prompt, context=None, temperature=0.2, max_new_tokens=512):
+    if context:
+        full_input = f"Context: {context}\n\nQuestion: {prompt}\nAnswer:"
+    else:
+        full_input = f"Question: {prompt}\nAnswer:"
+    inputs = tokenizer(full_input, return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=0.9,
+        do_sample=True
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# -----------------------------
+# 5️⃣  Pipeline function
+# -----------------------------
+def compliance_agent_pipeline(prompt, ticker=None, use_web=True):
+    tokenizer, model = load_model()
+    context = ""
+
+    # Optional: integrate external tools
+    if ticker:
+        fin_data = fetch_financial_data(ticker)
+        context += f"Market data for {ticker}: {json.dumps(fin_data, indent=2)}\n\n"
+
+    if use_web:
+        web_context = retrieve_context_from_web(prompt)
+        context += f"Retrieved context: {web_context}\n\n"
+
+    # Generate model output
+    response = generate_response(model, tokenizer, prompt, context)
+    return response
+
+
+# -----------------------------
+# 6️⃣  Example run
+# -----------------------------
+if __name__ == "__main__":
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    print("🧠 FinGPT Task II - Compliance Agent Inference")
+    query = input("Enter your compliance query: ")
+
+    # Example: "Summarize the compliance risk of Meta in 2025"
+    result = compliance_agent_pipeline(prompt=query, ticker="META", use_web=True)
+
+    print("\n=== AI Compliance Agent Response ===\n")
+    print(result)