Upload 2 files

app.py

@@ -1,16 +1,18 @@
 import streamlit as st
 import fitz  # PyMuPDF for PDF processing
 import pandas as pd
-from transformers import pipeline
 import os
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
 
 # Get the Hugging Face token from the environment variables
 hf_token = os.getenv("HF_API_TOKEN")
 
+
 # Load the model (Meta-Llama 3.1 8B)
 @st.cache_resource
 def load_model():
-    model = pipeline("text2text-generation", model="meta-llama/Meta-Llama-3.1-8B-Instruct", use_auth_token=hf_token)
+    model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct", use_auth_token=hf_token)
     return model
 
 model = load_model()
@@ -65,7 +67,10 @@ def generate_extraction_prompt(chunk):
 # Function to query Meta-Llama for each chunk
 def extract_financial_metrics_from_chunk(chunk):
     prompt = generate_extraction_prompt(chunk)
-    response = model(prompt)
+    model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct", use_auth_token=hf_token)
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct", use_auth_token=hf_token)
+    nlp = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    response = nlp(prompt)
     return response[0]['generated_text']
 
 # Process the PDF text through the model