Llama-Test / app.py
SagaMKM's picture
Update app.py
f231758 verified
raw
history blame
1.84 kB
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
huggingface-cli login
# Load LLaMA model and tokenizer from Hugging Face
model_name = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Load financial dataset to enrich responses
dataset = load_dataset("gbharti/finance-alpaca")
# Helper function to extract dataset info (optional enhancement)
def get_insight_from_dataset():
sample = dataset["train"].shuffle(seed=42).select([0])[0]
return f"Example insight: {sample['text']}"
# Function to process user input and generate financial advice
def financial_advisor(user_input):
# Tokenize the user input
inputs = tokenizer(user_input, return_tensors="pt")
# Generate response using the LLaMA model
outputs = model.generate(**inputs, max_length=256, num_return_sequences=1)
advice = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Get additional insight from dataset to enrich advice (optional)
insight = get_insight_from_dataset()
# Combine the advice and the insight
full_response = f"Advice: {advice}\n\n{insight}"
return full_response
# Create Gradio Interface
interface = gr.Interface(
fn=financial_advisor,
inputs=gr.Textbox(lines=5, placeholder="Enter your financial question..."),
outputs="text",
title="AI Financial Advisor",
description="Ask me anything related to finance, investments, savings, and more.",
examples=[
"Should I invest in stocks or real estate?",
"How can I save more money on a tight budget?",
"What are some good investment options for retirement?",
]
)
# Launch the Gradio app
if __name__ == "__main__":
interface.launch()