amiguel commited on
Commit
3358b89
Β·
verified Β·
1 Parent(s): d3fb5b0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -21
app.py CHANGED
@@ -2,24 +2,37 @@ import streamlit as st
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  from PyPDF2 import PdfReader
4
  import pandas as pd
 
 
 
 
 
5
 
6
  # Title and emojis
7
  st.title("πŸš€ WizNerd Insp πŸš€")
8
 
9
  # Sidebar for file uploads
10
  st.sidebar.header("Upload Files")
11
- uploaded_xlsx = st.sidebar.file_uploader("Upload XLSX File", type=["xlsx"])
12
- uploaded_pdf = st.sidebar.file_uploader("Upload PDF File", type=["pdf"])
13
 
14
  # Load the HuggingFace model and tokenizer
15
  @st.cache_resource
16
  def load_model():
17
  model_name = "amiguel/optimizedModelLinsting6.1"
18
- tokenizer = AutoTokenizer.from_pretrained(model_name)
19
- model = AutoModelForCausalLM.from_pretrained(model_name)
 
 
 
 
 
20
  return tokenizer, model
21
 
22
- tokenizer, model = load_model()
 
 
 
 
23
 
24
  # Prompt style
25
  prompt_style = """
@@ -37,14 +50,10 @@ Please answer the following inspection scope question.
37
 
38
  # Function to process user input and generate response
39
  def generate_response(input_text):
40
- # Format the input using the prompt style
41
  formatted_input = prompt_style.format(input_text, "", "")
42
-
43
- # Tokenize and generate response
44
  inputs = tokenizer(formatted_input, return_tensors="pt", truncation=True, max_length=512)
45
  outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True)
46
  response = tokenizer.decode(outputs[0], skip_special_tokens=True)
47
-
48
  return response
49
 
50
  # Main chat interface
@@ -57,15 +66,14 @@ if st.button("Submit"):
57
  st.write(response)
58
 
59
  # Process uploaded files
60
- if uploaded_xlsx:
61
- st.write("Processing XLSX file...")
62
- df = pd.read_excel(uploaded_xlsx)
63
- st.write(df)
64
-
65
- if uploaded_pdf:
66
- st.write("Processing PDF file...")
67
- pdf_reader = PdfReader(uploaded_pdf)
68
- text = ""
69
- for page in pdf_reader.pages:
70
- text += page.extract_text()
71
- st.write(text)
 
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  from PyPDF2 import PdfReader
4
  import pandas as pd
5
+ import os
6
+ from dotenv import load_dotenv
7
+
8
+ # Load environment variables
9
+ load_dotenv()
10
 
11
  # Title and emojis
12
  st.title("πŸš€ WizNerd Insp πŸš€")
13
 
14
  # Sidebar for file uploads
15
  st.sidebar.header("Upload Files")
16
+ uploaded_file = st.sidebar.file_uploader("Upload XLSX or PDF File", type=["xlsx", "pdf"])
 
17
 
18
  # Load the HuggingFace model and tokenizer
19
  @st.cache_resource
20
  def load_model():
21
  model_name = "amiguel/optimizedModelLinsting6.1"
22
+ hf_token = os.getenv("HUGGINGFACE_TOKEN") # Load token from .env
23
+ if hf_token:
24
+ tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
25
+ model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
26
+ else:
27
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
28
+ model = AutoModelForCausalLM.from_pretrained(model_name)
29
  return tokenizer, model
30
 
31
+ try:
32
+ tokenizer, model = load_model()
33
+ except Exception as e:
34
+ st.error(f"Error loading model: {e}")
35
+ st.info("Ensure the model name is correct or provide a valid Hugging Face token.")
36
 
37
  # Prompt style
38
  prompt_style = """
 
50
 
51
  # Function to process user input and generate response
52
  def generate_response(input_text):
 
53
  formatted_input = prompt_style.format(input_text, "", "")
 
 
54
  inputs = tokenizer(formatted_input, return_tensors="pt", truncation=True, max_length=512)
55
  outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True)
56
  response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
57
  return response
58
 
59
  # Main chat interface
 
66
  st.write(response)
67
 
68
  # Process uploaded files
69
+ if uploaded_file:
70
+ st.write(f"Processing {uploaded_file.type} file...")
71
+ if uploaded_file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
72
+ df = pd.read_excel(uploaded_file)
73
+ st.write(df)
74
+ elif uploaded_file.type == "application/pdf":
75
+ pdf_reader = PdfReader(uploaded_file)
76
+ text = ""
77
+ for page in pdf_reader.pages:
78
+ text += page.extract_text()
79
+ st.write(text)