Update app.py

app.py

@@ -11,23 +11,9 @@ from huggingface_hub import hf_hub_download, login
 import logging
 import pandas as pd
 import torch
-import yfinance as yf
-from datetime import datetime, timedelta
-from math import sqrt
-import time
-import base64
-import io
-import numpy as np
-try:
-    import scipy.optimize as opt
-except ModuleNotFoundError:
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "scipy"])
-    import scipy.optimize as opt
-
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-
 # Install llama-cpp-python with appropriate backend
 try:
     from llama_cpp import Llama
@@ -39,40 +25,34 @@ except ModuleNotFoundError:
     else:
         logger.info("Installing llama-cpp-python without additional flags.")
         subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python", "--force-reinstall", "--upgrade", "--no-cache-dir"])
-    from llama_cpp import Llama
-
+from llama_cpp import Llama
 # Install yfinance if not present (for CAGR calculations)
 try:
     import yfinance as yf
 except ModuleNotFoundError:
     subprocess.check_call([sys.executable, "-m", "pip", "install", "yfinance"])
-    import yfinance as yf
-
+import yfinance as yf
 # Import pandas for handling DataFrame column structures
 import pandas as pd
-
 # Additional imports for visualization and file handling
 try:
     import matplotlib.pyplot as plt
     from PIL import Image
     import io
 except ModuleNotFoundError:
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "matplotlib", "pillow", "numpy"])
-    import matplotlib.pyplot as plt
-    from PIL import Image
-    import io
-    import numpy as np
-
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "matplotlib", "pillow"])
+import matplotlib.pyplot as plt
+from PIL import Image
+import io
 MAX_MAX_NEW_TOKENS = 512
 DEFAULT_MAX_NEW_TOKENS = 512
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "1024"))
-
 DESCRIPTION = """# FinChat: Investing Q&A (Optimized for Speed)
-This application delivers an interactive chat interface powered by a highly efficient, small AI model adapted for addressing investing and finance inquiries through specialized prompt engineering. It ensures rapid, reasoned responses to user queries. Duplicate this Space for customization or queue-free deployment.<p>Running on CPU or GPU if available. Using Phi-2 model for faster inference. Inference is heavily optimized for responses in under 10 seconds for simple queries, with output limited to 250 tokens maximum. For longer responses, increase 'Max New Tokens' in Advanced Settings. Brief delays may occur in free-tier environments due to shared resources, but typical generation speeds are improved with the smaller model.</p>"""
-
-LICENSE = """<p/>---
+This application delivers an interactive chat interface powered by a highly efficient, small AI model adapted for addressing investing and finance inquiries through specialized prompt engineering. It ensures rapid, reasoned responses to user queries. Duplicate this Space for customization or queue-free deployment.
+<p>Running on CPU or GPU if available. Using Phi-2 model for faster inference. Inference is heavily optimized for responses in under 10 seconds for simple queries, with output limited to 250 tokens maximum. For longer responses, increase 'Max New Tokens' in Advanced Settings. Brief delays may occur in free-tier environments due to shared resources, but typical generation speeds are improved with the smaller model.</p>"""
+LICENSE = """<p/>
+---
 This application employs the Phi-2 model, governed by Microsoft's Terms of Use. Refer to the [model card](https://huggingface.co/TheBloke/phi-2-GGUF) for details."""
-
 # Load the model (skip fine-tuning for faster startup)
 try:
     model_path = hf_hub_download(
@@ -83,10 +63,10 @@ try:
     llm = Llama(
         model_path=model_path,
         n_ctx=1024,
-        n_batch=1024,  # Increased for faster processing
+        n_batch=1024, # Increased for faster processing
         n_threads=multiprocessing.cpu_count(),
         n_gpu_layers=n_gpu_layers,
-        chat_format="chatml"  # Phi-2 uses ChatML format in llama.cpp
+        chat_format="chatml" # Phi-2 uses ChatML format in llama.cpp
     )
     logger.info(f"Model loaded successfully with n_gpu_layers={n_gpu_layers}.")
     # Warm up the model for faster initial inference
@@ -95,361 +75,326 @@ try:
 except Exception as e:
     logger.error(f"Error loading model: {str(e)}")
     raise
-
 # Register explicit close for llm to avoid destructor error
 atexit.register(llm.close)
-
 DEFAULT_SYSTEM_PROMPT = """You are FinChat, a knowledgeable AI assistant specializing in investing and finance. Provide accurate, helpful, reasoned, and concise answers to investing questions. Always base responses on reliable information and advise users to consult professionals for personalized advice.
 Always respond exclusively in English. Use bullet points for clarity.
-Do not substitute or alter stock symbols provided in the user's query. Always use the exact tickers mentioned.
 Example:
-User: average return for AAPL between 2010 and 2020
+User: average return for TSLA between 2010 and 2020
 Assistant:
-- AAPL CAGR (2010-2020): ~27.24%
+- TSLA CAGR (2010-2020): ~63.01%
 - Represents average annual return with compounding
-- Past performance is not indicative of future results."""
-
-logs = []
-
-# Function to calculate CAGR using yfinance
-def calculate_cagr(ticker, start_date, end_date):
-    try:
-        data = yf.download(ticker, start=start_date, end=end_date)
-        if data.empty:
-            return None
-        start_price = data['Adj Close'].iloc[0]
-        end_price = data['Adj Close'].iloc[-1]
-        num_years = (data.index[-1] - data.index[0]).days / 365.25
-        cagr = (end_price / start_price) ** (1 / num_years) - 1
-        return cagr * 100  # Return as percentage
-    except Exception as e:
-        logger.error(f"Error calculating CAGR for {ticker}: {str(e)}")
-        return None
-
-# Function to calculate risk metrics using yfinance
-def calculate_risk_metrics(ticker, years=5):
-    try:
-        end_date = datetime.now().strftime('%Y-%m-%d')
-        start_date = (datetime.now() - timedelta(days=365 * years)).strftime('%Y-%m-%d')
-        data = yf.download(ticker, start=start_date, end=end_date)
-        if data.empty:
-            return None, None
-        returns = data['Adj Close'].pct_change().dropna()
-        volatility = returns.std() * sqrt(252) * 100  # Annualized volatility in percent
-        mean_return = returns.mean() * 252  # Annualized mean return
-        risk_free_rate = 0.02  # Assumed risk-free rate (e.g., 2%)
-        sharpe = (mean_return - risk_free_rate) / (volatility / 100)  # Sharpe ratio
-        return volatility, sharpe
-    except Exception as e:
-        logger.error(f"Error calculating risk metrics for {ticker}: {str(e)}")
-        return None, None
-
-# Function for inline plot
-def generate_plot(ticker, period='5y'):
-    try:
-        data = yf.download(ticker, period=period)
-        if data.empty:
-            return "Unable to fetch data for plotting."
-        plt.figure(figsize=(10, 5))
-        plt.plot(data['Adj Close'], label='Adjusted Close')
-        plt.title(f'{ticker} Price History ({period})')
-        plt.xlabel('Date')
-        plt.ylabel('Price (USD)')
-        plt.legend()
-        plt.grid(True)
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight')
-        buf.seek(0)
-        b64 = base64.b64encode(buf.read()).decode('utf-8')
-        plt.close()
-        return f"![{ticker} Price Chart](data:image/png;base64,{b64})"
-    except Exception as e:
-        logger.error(f"Error generating plot for {ticker}: {str(e)}")
-        return "Error generating plot."
-
-# Function for portfolio optimization using scipy
-def portfolio_optimization(tickers, target_return=None):
-    try:
-        data = yf.download(tickers, period='5y')['Adj Close']
-        returns = data.pct_change().dropna()
-        mean_returns = returns.mean() * 252
-        cov_matrix = returns.cov() * 252
-        num_assets = len(tickers)
-
-        def portfolio_volatility(weights):
-            return np.sqrt(np.dot(weights.T, np.dot(cov_matrix, weights)))
-
-        constraints = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})
-        bounds = tuple((0, 1) for _ in range(num_assets))
-        initial_guess = np.array(num_assets * [1. / num_assets])
-
-        if target_return:
-            # Maximize Sharpe or min vol for target return
-            def objective(weights):
-                ret = np.sum(mean_returns * weights)
-                vol = portfolio_volatility(weights)
-                return - (ret - 0.02) / vol if vol != 0 else np.inf  # Neg Sharpe
-            cons = [{'type': 'eq', 'fun': lambda x: np.sum(x) - 1},
-                    {'type': 'eq', 'fun': lambda x: np.sum(mean_returns * x) - target_return}]
-            result = opt.minimize(objective, initial_guess, method='SLSQP', bounds=bounds, constraints=cons)
-        else:
-            # Minimize volatility
-            result = opt.minimize(portfolio_volatility, initial_guess, method='SLSQP',
-                                  bounds=bounds, constraints=constraints)
-
-        if result.success:
-            weights = dict(zip(tickers, result.x))
-            return weights
-        else:
-            return {ticker: 1/len(tickers) for ticker in tickers}  # Fallback equal weights
-    except Exception as e:
-        logger.error(f"Error in portfolio optimization: {str(e)}")
-        return {ticker: 1/len(tickers) for ticker in tickers}
-
-# Assuming the generate function handles the chat logic (extended to include risk comparison)
+- Past performance not indicative of future results
+- Consult a financial advisor"""
+# Company name to ticker mapping (expand as needed)
+COMPANY_TO_TICKER = {
+    "opendoor": "OPEN",
+    "tesla": "TSLA",
+    "apple": "AAPL",
+    "amazon": "AMZN",
+    # Add more mappings for common companies
+}
 def generate(
     message: str,
-    history: list[tuple[str, str]],
-    system_prompt: str,
-    max_new_tokens: int,
-    temperature: float,
-    top_p: float,
-    top_k: int,
-    logs_state: list
-) -> tuple[Iterator[str], list]:
-    start_time = time.time()
-    if not system_prompt:
-        system_prompt = DEFAULT_SYSTEM_PROMPT
-
-    full_response = ""
-    # Detect CAGR query
-    cagr_match = re.search(r'average return for (\w+) between (\d{4}) and (\d{4})', message.lower())
-    if cagr_match:
-        ticker = cagr_match.group(1).upper()
-        start_year = cagr_match.group(2)
-        end_year = cagr_match.group(3)
-        start_date = f"{start_year}-01-01"
-        end_date = f"{end_year}-12-31"
-        cagr = calculate_cagr(ticker, start_date, end_date)
-        if cagr is not None:
-            response = f"- {ticker} CAGR ({start_year}-{end_year}): ~{cagr:.2f}%\n- Represents average annual return with compounding\n- Past performance is not indicative of future results.\n- Consult a financial advisor for personalized advice."
-            yield response
-            full_response = response
-        else:
-            response = "Unable to calculate CAGR for the specified period."
-            yield response
-            full_response = response
-        end_time = time.time()
-        logs_state.append({
-            'timestamp': datetime.now().isoformat(),
-            'query': message,
-            'response': full_response,
-            'response_length': len(full_response.split()),
-            'generation_time': end_time - start_time,
-            'token_efficiency': len(full_response.split()) / max_new_tokens
-        })
-        return iter([]), logs_state  # No more yield
-
-    # Detect risk comparison query
-    risk_match = re.search(r'which stock is riskier (\w+) or (\w+)', message.lower())
-    if risk_match:
-        ticker1 = risk_match.group(1).upper()
-        ticker2 = risk_match.group(2).upper()
-        vol1, sharpe1 = calculate_risk_metrics(ticker1)
-        vol2, sharpe2 = calculate_risk_metrics(ticker2)
-        if vol1 is None or vol2 is None:
-            response = "Unable to fetch risk metrics for one or both tickers."
-            yield response
-            full_response = response
-        else:
-            if vol1 > vol2:
-                riskier = ticker1
-                less_risky = ticker2
-                higher_vol = vol1
-                lower_vol = vol2
-                riskier_sharpe = sharpe1
-                less_sharpe = sharpe2
-            else:
-                riskier = ticker2
-                less_risky = ticker1
-                higher_vol = vol2
-                lower_vol = vol1
-                riskier_sharpe = sharpe2
-                less_sharpe = sharpe1
-            response = f"- {riskier} is riskier compared to {less_risky}.\n- It has a higher annualized standard deviation ({higher_vol:.2f}% vs {lower_vol:.2f}%) and a lower Sharpe ratio ({riskier_sharpe:.2f} vs {less_sharpe:.2f}), indicating greater volatility and potentially lower risk-adjusted returns.\n- Calculations based on the past 5 years of data.\n- Past performance is not indicative of future results. Consult a financial advisor for personalized advice."
+    chat_history: list[dict],
+    system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
+    temperature: float = 0.6,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.2,
+) -> Iterator[str]:
+    logger.info(f"Generating response for message: {message}")
+    lower_message = message.lower().strip()
+    if lower_message in ["hi", "hello"]:
+        response = "I'm FinChat, your financial advisor. Ask me anything finance-related!"
+        logger.info("Quick response for 'hi'/'hello' generated.")
+        yield response
+        return
+    if "what is cagr" in lower_message:
+        response = """- CAGR stands for Compound Annual Growth Rate.
+- It measures the mean annual growth rate of an investment over a specified period longer than one year, accounting for compounding.
+- Formula: CAGR = (Ending Value / Beginning Value)^(1 / Number of Years) - 1
+- Useful for comparing investments over time.
+- Past performance not indicative of future results. Consult a financial advisor."""
+        logger.info("Quick response for 'what is cagr' generated.")
+        yield response
+        return
+    # Check for compound interest queries
+    compound_match = re.search(r'(?:save|invest|deposit)\s*\$?([\d,]+(?:\.\d+)?)\s*(?:right now|today)?\s*(?:under|at)\s*([\d.]+)%\s*(?:interest|rate)?\s*(?:annually|per year)?\s*over\s*(\d+)\s*years?', lower_message)
+    if compound_match:
+        try:
+            principal_str = compound_match.group(1).replace(',', '')
+            principal = float(principal_str)
+            rate = float(compound_match.group(2)) / 100
+            years = int(compound_match.group(3))
+            if principal <= 0 or rate < 0 or years <= 0:
+                yield "Invalid input values: principal, rate, and years must be positive."
+                return
+            balance = principal * (1 + rate) ** years
+            response = (
+                f"- Starting with ${principal:,.2f} at {rate*100:.2f}% annual interest, compounded annually over {years} years.\n"
+                f"- Projected balance in year {years}: ${balance:,.2f}\n"
+                f"- Assumptions: Annual compounding; no additional deposits or withdrawals.\n"
+                f"- This is a projection; actual results may vary. Consult a financial advisor."
+            )
+            logger.info("Compound interest response generated.")
             yield response
-            full_response = response
-        end_time = time.time()
-        logs_state.append({
-            'timestamp': datetime.now().isoformat(),
-            'query': message,
-            'response': full_response,
-            'response_length': len(full_response.split()),
-            'generation_time': end_time - start_time,
-            'token_efficiency': len(full_response.split()) / max_new_tokens
-        })
-        return iter([]), logs_state
-
-    # Detect plot/chart query
-    plot_match = re.search(r'(plot|chart)\s+(\w+)(?:\s+(historical|price|volatility))?', message.lower())
-    if plot_match:
-        ticker = plot_match.group(2).upper()
-        plot_type = plot_match.group(3) if plot_match.group(3) else 'price'
-        if plot_type == 'volatility':
-            # Simple volatility plot (returns histogram)
+            return
+        except ValueError as ve:
+            logger.error(f"Error parsing compound interest query: {str(ve)}")
+            yield "Error parsing query: Please ensure amount, rate, and years are valid numbers."
+            return
+    # Check for CAGR/average return queries (use re.search for flexible matching)
+    match = re.search(r'(?:average return|cagr) for ([\w\s,]+(?:and [\w\s,]+)?) between (\d{4}) and (\d{4})', lower_message)
+    if match:
+        tickers_str, start_year, end_year = match.groups()
+        tickers = [t.strip().upper() for t in re.split(r',|\band\b', tickers_str) if t.strip()]
+        # Apply company-to-ticker mapping
+        for i in range(len(tickers)):
+            lower_ticker = tickers[i].lower()
+            if lower_ticker in COMPANY_TO_TICKER:
+                tickers[i] = COMPANY_TO_TICKER[lower_ticker]
+        responses = []
+        if int(end_year) <= int(start_year):
+            yield "The specified time period is invalid (end year must be after start year)."
+            return
+        for ticker in tickers:
             try:
-                data = yf.download(ticker, period='1y')
-                returns = data['Adj Close'].pct_change().dropna()
-                plt.figure(figsize=(10, 5))
-                plt.hist(returns, bins=50, alpha=0.7)
-                plt.title(f'{ticker} Daily Returns Distribution (1Y)')
-                plt.xlabel('Return')
-                plt.ylabel('Frequency')
-            except:
-                plot_type = 'price'  # Fallback
-        if plot_type != 'volatility':
-            plot_md = generate_plot(ticker)
-            response = f"Price chart for {ticker}:\n{plot_md}\n- This visualizes the historical adjusted close prices.\n- Past performance is not indicative of future results. Consult a financial advisor."
-            yield response
-            full_response = response
-        else:
-            # For volatility, similar
-            buf = io.BytesIO()
-            plt.savefig(buf, format='png', bbox_inches='tight')
-            buf.seek(0)
-            b64 = base64.b64encode(buf.read()).decode('utf-8')
-            plt.close()
-            plot_md = f"![{ticker} Volatility](data:image/png;base64,{b64})"
-            response = f"Volatility chart for {ticker}:\n{plot_md}\n- Histogram of daily returns over the past year."
-            yield response
-            full_response = response
-        end_time = time.time()
-        logs_state.append({
-            'timestamp': datetime.now().isoformat(),
-            'query': message,
-            'response': full_response,
-            'response_length': len(full_response.split()),
-            'generation_time': end_time - start_time,
-            'token_efficiency': len(full_response.split()) / max_new_tokens
-        })
-        return iter([]), logs_state
-
-    # Detect portfolio optimization query
-    port_match = re.search(r'optimize\s+portfolio\s+for\s+([\w,\s]+)(?:\s+with\s+(risk|return)\s+tolerance\s+([\d.]+))?', message.lower())
-    if port_match:
-        tickers_str = port_match.group(1).strip()
-        tickers = [t.strip().upper() for t in re.split(r'[,;]', tickers_str) if t.strip()]
-        target = None
-        if port_match.group(3):
-            target = float(port_match.group(3))
-            if port_match.group(2) == 'risk':
-                # For risk tolerance, min vol with vol <= target (but simplify to min vol)
-                pass  # Use default min vol
-            else:
-                target_return = target
-        weights = portfolio_optimization(tickers, target_return=target if 'return' in (port_match.group(2) or '') else None)
-        df = pd.DataFrame(list(weights.items()), columns=['Ticker', 'Weight'])
-        df['Weight'] = df['Weight'].round(4)
-        table_md = df.to_markdown(index=False)
-        response = f"- Suggested portfolio weights for {', '.join(tickers)}:\n{table_md}\n- Based on minimum variance optimization (or target return if specified).\n- Assumes 5-year historical data for means and covariances.\n- Past performance is not indicative of future results. Consult a financial advisor for personalized advice."
-        yield response
-        full_response = response
-        end_time = time.time()
-        logs_state.append({
-            'timestamp': datetime.now().isoformat(),
-            'query': message,
-            'response': full_response,
-            'response_length': len(full_response.split()),
-            'generation_time': end_time - start_time,
-            'token_efficiency': len(full_response.split()) / max_new_tokens
-        })
-        return iter([]), logs_state
-
-    # For other queries, fall back to LLM generation
+                # Download data with adjusted close prices
+                data = yf.download(ticker, start=f"{start_year}-01-01", end=f"{end_year}-12-31", progress=False, auto_adjust=False)
+                # Handle potential MultiIndex columns in newer yfinance versions
+                if isinstance(data.columns, pd.MultiIndex):
+                    data.columns = data.columns.droplevel(1)
+                if not data.empty:
+                    # Check if 'Adj Close' column exists
+                    if 'Adj Close' not in data.columns:
+                        responses.append(f"- {ticker}: Error - Adjusted Close price data not available.")
+                        logger.error(f"No 'Adj Close' column for {ticker}.")
+                        continue
+                    # Ensure data is not MultiIndex for single ticker (already handled)
+                    initial = data['Adj Close'].iloc[0]
+                    final = data['Adj Close'].iloc[-1]
+                    start_date = data.index[0]
+                    end_date = data.index[-1]
+                    days = (end_date - start_date).days
+                    years = days / 365.25
+                    if years > 0 and pd.notna(initial) and pd.notna(final):
+                        cagr = ((final / initial) ** (1 / years) - 1) * 100
+                        responses.append(f"- {ticker}: ~{cagr:.2f}%")
+                    else:
+                        responses.append(f"- {ticker}: Invalid period or missing price data.")
+                else:
+                    responses.append(f"- {ticker}: No historical data available between {start_year} and {end_year}.")
+            except Exception as e:
+                logger.error(f"Error calculating CAGR for {ticker}: {str(e)}")
+                responses.append(f"- {ticker}: Error calculating CAGR - {str(e)}")
+        full_response = f"CAGR for the requested stocks from {start_year} to {end_year}:\n" + "\n".join(responses) + "\n- Represents average annual returns with compounding\n- Past performance not indicative of future results\n- Consult a financial advisor"
+        full_response = re.sub(r'<\|(?:im_start|im_end|system|user|assistant)\|>|</s>|\[END\]', '', full_response).strip() # Clean any trailing tokens
+        # Estimate token count to ensure response fits within max_new_tokens
+        response_tokens = len(llm.tokenize(full_response.encode("utf-8"), add_bos=False))
+        if response_tokens > max_new_tokens:
+            logger.warning(f"CAGR response tokens ({response_tokens}) exceed max_new_tokens ({max_new_tokens}). Truncating to first complete sentence.")
+            sentence_endings = ['.', '!', '?']
+            first_sentence_end = min([full_response.find(ending) + 1 for ending in sentence_endings if full_response.find(ending) != -1], default=len(full_response))
+            full_response = full_response[:first_sentence_end] if first_sentence_end > 0 else "Response truncated due to length; please increase Max New Tokens."
+        logger.info("CAGR response generated.")
+        yield full_response
+        return
+    # Build conversation messages (limit history to last 3 for speed)
     conversation = [{"role": "system", "content": system_prompt}]
-    for user, assistant in history:
-        conversation.extend([
-            {"role": "user", "content": user},
-            {"role": "assistant", "content": assistant}
-        ])
+    for msg in chat_history[-3:]: # Reduced from 5 to 3 for faster processing
+        if msg["role"] == "user":
+            conversation.append({"role": "user", "content": msg["content"]})
+        elif msg["role"] == "assistant":
+            conversation.append({"role": "assistant", "content": msg["content"]})
     conversation.append({"role": "user", "content": message})
-
-    # Generate response using LLM (streamed)
-    response = llm.create_chat_completion(
-        messages=conversation,
-        max_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        stream=True
+    # Approximate token length check and truncate if necessary
+    prompt_text = "\n".join(d["content"] for d in conversation)
+    input_tokens = llm.tokenize(prompt_text.encode("utf-8"), add_bos=False)
+    while len(input_tokens) > MAX_INPUT_TOKEN_LENGTH:
+        logger.warning(f"Input tokens ({len(input_tokens)}) exceed limit ({MAX_INPUT_TOKEN_LENGTH}). Truncating history.")
+        if len(conversation) > 2: # Preserve system prompt and current user message
+            conversation.pop(1) # Remove oldest user/assistant pair
+            prompt_text = "\n".join(d["content"] for d in conversation)
+            input_tokens = llm.tokenize(prompt_text.encode("utf-8"), add_bos=False)
+        else:
+            yield "Error: Input is too long even after truncation. Please shorten your query."
+            return
+    # Generate response with sentence boundary checking and token cleanup
+    try:
+        response = ""
+        sentence_buffer = ""
+        token_count = 0
+        stream = llm.create_chat_completion(
+            messages=conversation,
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repeat_penalty=repetition_penalty,
+            stream=True
+        )
+        sentence_endings = ['.', '!', '?']
+        for chunk in stream:
+            delta = chunk["choices"][0]["delta"]
+            if "content" in delta and delta["content"] is not None:
+                # Clean the chunk by removing ChatML tokens or similar
+                cleaned_chunk = re.sub(r'<\|(?:im_start|im_end|system|user|assistant)\|>|</s>|\[END\]', '', delta["content"])
+                if not cleaned_chunk:
+                    continue
+                sentence_buffer += cleaned_chunk
+                response += cleaned_chunk
+                # Approximate token count for the chunk
+                chunk_tokens = len(llm.tokenize(cleaned_chunk.encode("utf-8"), add_bos=False))
+                token_count += chunk_tokens
+                # Check for sentence boundary
+                if any(sentence_buffer.strip().endswith(ending) for ending in sentence_endings):
+                    yield response
+                    sentence_buffer = "" # Clear buffer after yielding a complete sentence
+                # Removed early truncation to allow full token utilization
+            if chunk["choices"][0]["finish_reason"] is not None:
+                # Yield any remaining complete sentence in the buffer
+                if sentence_buffer.strip():
+                    last_sentence_end = max([sentence_buffer.rfind(ending) for ending in sentence_endings if sentence_buffer.rfind(ending) != -1], default=-1)
+                    if last_sentence_end != -1:
+                        response = response[:response.rfind(sentence_buffer) + last_sentence_end + 1]
+                        yield response
+                    else:
+                        yield response
+                else:
+                    yield response
+                break
+        logger.info("Response generation completed.")
+    except ValueError as ve:
+        if "exceed context window" in str(ve):
+            yield "Error: Prompt too long for context window. Please try a shorter query or clear history."
+        else:
+            logger.error(f"Error during response generation: {str(ve)}")
+            yield f"Error generating response: {str(ve)}"
+    except Exception as e:
+        logger.error(f"Error during response generation: {str(e)}")
+        yield f"Error generating response: {str(e)}"
+def process_portfolio(df, growth_rate):
+    if df is None or len(df) == 0:
+        return "", None
+    # Convert to DataFrame if needed
+    if not isinstance(df, pd.DataFrame):
+        df = pd.DataFrame(df, columns=["Ticker", "Shares", "Avg Cost", "Current Price"])
+    df = df.dropna(subset=["Ticker"])
+    portfolio = {}
+    for _, row in df.iterrows():
+        ticker = row["Ticker"].upper() if pd.notna(row["Ticker"]) else None
+        if not ticker:
+            continue
+        shares = float(row["Shares"]) if pd.notna(row["Shares"]) else 0
+        cost = float(row["Avg Cost"]) if pd.notna(row["Avg Cost"]) else 0
+        price = float(row["Current Price"]) if pd.notna(row["Current Price"]) else 0
+        value = shares * price
+        portfolio[ticker] = {'shares': shares, 'cost': cost, 'price': price, 'value': value}
+    if not portfolio:
+        return "", None
+    total_value_now = sum(v['value'] for v in portfolio.values())
+    allocations = {k: v['value'] / total_value_now for k, v in portfolio.items()} if total_value_now > 0 else {}
+    fig_alloc, ax_alloc = plt.subplots()
+    ax_alloc.pie(allocations.values(), labels=allocations.keys(), autopct='%1.1f%%')
+    ax_alloc.set_title('Portfolio Allocation')
+    buf_alloc = io.BytesIO()
+    fig_alloc.savefig(buf_alloc, format='png')
+    buf_alloc.seek(0)
+    chart_alloc = Image.open(buf_alloc)
+    plt.close(fig_alloc) # Close the figure to free memory
+    def project_value(value, years, rate):
+        return value * (1 + rate / 100) ** years
+    total_value_1yr = sum(project_value(v['value'], 1, growth_rate) for v in portfolio.values())
+    total_value_2yr = sum(project_value(v['value'], 2, growth_rate) for v in portfolio.values())
+    total_value_5yr = sum(project_value(v['value'], 5, growth_rate) for v in portfolio.values())
+    total_value_10yr = sum(project_value(v['value'], 10, growth_rate) for v in portfolio.values())
+    data_str = (
+        "User portfolio:\n" +
+        "\n".join(f"- {k}: {v['shares']} shares, avg cost {v['cost']}, current price {v['price']}, value ${v['value']:,.2f}" for k, v in portfolio.items()) +
+        f"\nTotal value now: ${total_value_now:,.2f}\nProjected (at {growth_rate}% annual growth):\n" +
+        f"- 1 year: ${total_value_1yr:,.2f}\n- 2 years: ${total_value_2yr:,.2f}\n- 5 years: ${total_value_5yr:,.2f}\n- 10 years: ${total_value_10yr:,.2f}"
     )
-    partial_text = ""
-    for chunk in response:
-        if "content" in chunk["choices"][0]["delta"]:
-            partial_text += chunk["choices"][0]["delta"]["content"]
-            yield partial_text
-            full_response = partial_text
-    end_time = time.time()
-    logs_state.append({
-        'timestamp': datetime.now().isoformat(),
-        'query': message,
-        'response': full_response,
-        'response_length': len(full_response.split()),
-        'generation_time': end_time - start_time,
-        'token_efficiency': len(full_response.split()) / max_new_tokens
-    })
-    return iter([]), logs_state
-
-def update_logs(logs_state):
-    if logs_state:
-        df = pd.DataFrame(logs_state)
+    return data_str, chart_alloc
+def fetch_current_prices(df):
+    if df is None or len(df) == 0:
         return df
-    return pd.DataFrame()
-
+    # Convert to DataFrame if needed
+    if not isinstance(df, pd.DataFrame):
+        df = pd.DataFrame(df, columns=["Ticker", "Shares", "Avg Cost", "Current Price"])
+    for i in df.index:
+        ticker = df.at[i, "Ticker"]
+        if pd.notna(ticker) and ticker.strip():
+            try:
+                price = yf.Ticker(ticker.upper()).info.get('currentPrice', None)
+                if price is not None:
+                    df.at[i, "Current Price"] = price
+            except Exception as e:
+                logger.warning(f"Failed to fetch price for {ticker}: {str(e)}")
+    return df
 # Gradio interface setup
-with gr.Blocks(theme=themes.Default()) as demo:
+with gr.Blocks(theme=themes.Soft(), css="""#chatbot {height: 800px; overflow: auto;}""") as demo:
     gr.Markdown(DESCRIPTION)
+    chatbot = gr.Chatbot(label="FinChat", type="messages")
+    msg = gr.Textbox(label="Ask a finance question", placeholder="e.g., 'What is CAGR?' or 'Average return for AAPL between 2010 and 2020'", info="Enter your query here. Portfolio data will be appended if provided.")
+    with gr.Row():
+        submit = gr.Button("Submit", variant="primary")
+        clear = gr.Button("Clear")
+    gr.Examples(
+        examples=["What is CAGR?", "Average return for AAPL between 2010 and 2020", "Hi", "Explain compound interest"],
+        inputs=msg,
+        label="Example Queries"
+    )
+    with gr.Accordion("Enter Portfolio for Projections", open=False):
+        portfolio_df = gr.Dataframe(
+            headers=["Ticker", "Shares", "Avg Cost", "Current Price"],
+            datatype=["str", "number", "number", "number"],
+            row_count=3,
+            col_count=(4, "fixed"),
+            label="Portfolio Data",
+            interactive=True
+        )
+        gr.Markdown("Enter your stocks here. You can add more rows by editing the table.")
+        fetch_button = gr.Button("Fetch Current Prices", variant="secondary")
+        fetch_button.click(fetch_current_prices, inputs=portfolio_df, outputs=portfolio_df)
+        growth_rate = gr.Slider(minimum=5, maximum=50, step=5, value=10, label="Annual Growth Rate (%)", interactive=True, info="Select the assumed annual growth rate for projections.")
+        growth_rate_label = gr.Markdown("**Selected Growth Rate: 10%**")
+    with gr.Accordion("Advanced Settings", open=False):
+        system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=6, info="Customize the AI's system prompt.")
+        temperature = gr.Slider(label="Temperature", value=0.6, minimum=0.0, maximum=1.0, step=0.05, info="Controls randomness: lower is more deterministic.")
+        top_p = gr.Slider(label="Top P", value=0.9, minimum=0.0, maximum=1.0, step=0.05, info="Nucleus sampling: higher includes more diverse tokens.")
+        top_k = gr.Slider(label="Top K", value=50, minimum=1, maximum=100, step=1, info="Top-K sampling: limits to top K tokens.")
+        repetition_penalty = gr.Slider(label="Repetition Penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05, info="Penalizes repeated tokens.")
+        max_new_tokens = gr.Slider(label="Max New Tokens", value=DEFAULT_MAX_NEW_TOKENS, minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, info="Maximum length of generated response.")
     gr.Markdown(LICENSE)
-    
-    with gr.Tabs():
-        with gr.TabItem("Chat"):
-            chatbot = gr.Chatbot()
-            msg = gr.Textbox(label="Enter your question")
-            with gr.Row():
-                submit = gr.Button("Submit")
-                clear = gr.Button("Clear")
-            advanced = gr.Accordion("Advanced Settings", open=False)
-            with advanced:
-                system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=6)
-                max_new_tokens = gr.Slider(minimum=1, maximum=MAX_MAX_NEW_TOKENS, value=DEFAULT_MAX_NEW_TOKENS, step=1, label="Max New Tokens")
-                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, label="Temperature")
-                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, step=0.1, label="Top P")
-                top_k = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top K")
-        
-        with gr.TabItem("Metrics"):
-            metrics_df = gr.Dataframe(headers=['timestamp', 'query', 'response', 'response_length', 'generation_time', 'token_efficiency'])
-    
-    logs_state = gr.State(logs)
-    
-    def submit_fn(msg, history, system_prompt, max_new_tokens, temperature, top_p, top_k, logs_state):
-        gen, new_logs = generate(msg, history, system_prompt, max_new_tokens, temperature, top_p, top_k, logs_state)
-        history.append((msg, ""))
-        for partial in gen:
-            history[-1] = (history[-1][0], partial)
-            yield history, "", new_logs
-        return history, "", new_logs
-    
-    submit.click(
-        submit_fn,
-        inputs=[msg, chatbot, system_prompt, max_new_tokens, temperature, top_p, top_k, logs_state],
-        outputs=[chatbot, msg, logs_state],
-        queue=False
-    ).then(
-        update_logs,
-        inputs=[logs_state],
-        outputs=[metrics_df]
+    def update_growth_rate_label(growth_rate):
+        return f"**Selected Growth Rate: {growth_rate}%**"
+    def user(message, history):
+        if not message:
+            return "", history
+        return "", history + [{"role": "user", "content": message}]
+    def bot(history, sys_prompt, temp, tp, tk, rp, mnt, portfolio_df, growth_rate):
+        if not history:
+            logger.warning("History is empty, initializing with user message.")
+            history = [{"role": "user", "content": ""}]
+        message = history[-1]["content"]
+        portfolio_data, chart_alloc = process_portfolio(portfolio_df, growth_rate)
+        message += "\n" + portfolio_data
+        history[-1]["content"] = message
+        history.append({"role": "assistant", "content": ""})
+        for new_text in generate(message, history[:-1], sys_prompt, mnt, temp, tp, tk, rp):
+            history[-1]["content"] = new_text
+            yield history, f"**Selected Growth Rate: {growth_rate}%**"
+        if chart_alloc:
+            history.append({"role": "assistant", "content": "", "image": chart_alloc})
+            yield history, f"**Selected Growth Rate: {growth_rate}%**"
+    growth_rate.change(update_growth_rate_label, inputs=growth_rate, outputs=growth_rate_label)
+    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, system_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens, portfolio_df, growth_rate], [chatbot, growth_rate_label]
+    )
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, system_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens, portfolio_df, growth_rate], [chatbot, growth_rate_label]
     )
-    
-    clear.click(lambda: ([], []), None, (chatbot, logs_state))
-
-demo.launch()
+    clear.click(lambda: [], None, chatbot, queue=False)
+demo.queue(max_size=128).launch()