from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import torch
# Alternative memory-efficient loading options without bitsandbytes
model_id = "tarun7r/Finance-Llama-8B"
print("Loading model with memory optimizations...")
# Option 1: Use FP16 (half precision) - reduces memory by ~50%
try:
print("Trying FP16 loading...")
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16, # Half precision
device_map="auto", # Automatic device placement
low_cpu_mem_usage=True, # Efficient CPU memory usage during loading
trust_remote_code=True
)
print("✓ Model loaded with FP16")
except Exception as e:
print(f"FP16 loading failed: {e}")
# Option 2: CPU offloading - some layers on GPU, some on CPU
try:
print("Trying CPU offloading...")
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="balanced", # Balance between GPU and CPU
low_cpu_mem_usage=True,
trust_remote_code=True
)
print("✓ Model loaded with CPU offloading")
except Exception as e:
print(f"CPU offloading failed: {e}")
# Option 3: Full CPU loading as fallback
print("Loading on CPU...")
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="cpu",
low_cpu_mem_usage=True,
trust_remote_code=True
)
print("✓ Model loaded on CPU")
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# Create pipeline
generator = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer
)
print("✓ Pipeline created successfully!")
# Your existing prompt code
finance_prompt_template = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction:
{}
### Input:
{}
### Response:
"""
# Update the system prompt to provide a more detailed description of the chatbot's role
messages = [
{"role": "system", "content": "You are a highly knowledgeable finance chatbot. Your purpose is to provide accurate, insightful, and actionable financial advice to users, tailored to their specific needs and contexts."},
{"role": "user", "content": "What strategies can an individual investor use to diversify their portfolio effectively in a volatile market?"},
]
# Update the generator call to use the messages
prompt = "\n".join([f"{msg['role'].capitalize()}: {msg['content']}" for msg in messages])
print("\n--- Generating Response ---")
try:
outputs = generator(
prompt,
#max_new_tokens=250, # Reduced for memory efficiency
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.eos_token_id,
# Memory efficient generation settings
num_beams=1, # No beam search to save memory
early_stopping=True,
use_cache=True
)
# Extract response
generated_text = outputs[0]['generated_text']
response_start = generated_text.rfind("### Response:")
if response_start != -1:
response = generated_text[response_start + len("### Response:"):].strip()
print("\n--- Response ---")
print(response)
else:
print(generated_text)
# Clean up GPU memory after generation
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception as e:
print(f"Generation error: {e}")