ryanpdwyer commited on
Commit
180f4e1
1 Parent(s): e424b8e

Switched to pipeline api

Browse files
Files changed (1) hide show
  1. app.py +19 -15
app.py CHANGED
@@ -1,5 +1,5 @@
1
  import streamlit as st
2
- from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
4
  import os
5
 
@@ -11,15 +11,19 @@ if not hf_token:
11
  st.error("Hugging Face token not found. Please add your HF_TOKEN to the Space secrets.")
12
  st.stop()
13
 
14
- # Load models and tokenizers
15
- @st.cache_resource
16
- def load_model_and_tokenizer(model_name):
17
- tokenizer = AutoTokenizer.from_pretrained(model_name)
18
- model = AutoModelForCausalLM.from_pretrained(model_name)
19
- return model, tokenizer
20
 
21
- model_8b, tokenizer_8b = load_model_and_tokenizer("unsloth/Meta-Llama-3.1-8B-bnb-4bit")
22
- model_8b_instruct, tokenizer_8b_instruct = load_model_and_tokenizer("SanctumAI/Meta-Llama-3.1-8B-Instruct-GGUF")
 
 
 
 
 
 
 
 
 
 
23
 
24
  def generate_text(model, tokenizer, prompt, max_length=100):
25
  inputs = tokenizer(prompt, return_tensors="pt")
@@ -36,14 +40,14 @@ if st.button("Generate"):
36
  if prompt:
37
  col1, col2 = st.columns(2)
38
 
39
- # with col1:
40
- # st.subheader("LLaMA-3.1-8B Output")
41
- # output_8b = generate_text(model_8b, tokenizer_8b, prompt, max_length)
42
- # st.write(output_8b)
43
 
44
  with col2:
45
  st.subheader("LLaMA-3.1-8B-Instruct Output")
46
- output_8b_instruct = generate_text(model_8b_instruct, tokenizer_8b_instruct, prompt, max_length)
47
- st.write(output_8b_instruct)
48
  else:
49
  st.warning("Please enter a prompt.")
 
1
  import streamlit as st
2
+ from transformers import pipeline
3
  import torch
4
  import os
5
 
 
11
  st.error("Hugging Face token not found. Please add your HF_TOKEN to the Space secrets.")
12
  st.stop()
13
 
 
 
 
 
 
 
14
 
15
+ @st.cache_resource
16
+ def load_pipeline(model_name):
17
+ with st.spinner(f'Loading {model_name}... This may take several minutes.'):
18
+ try:
19
+ pipe = pipeline("text-generation", model=model_name)
20
+ except Exception as e:
21
+ st.error(f"An error occurred: {e}")
22
+ st.stop()
23
+ return pipe
24
+
25
+ pipe8 = load_pipeline("unsloth/Meta-Llama-3.1-8B-bnb-4bit")
26
+ pipe8instruct = load_pipeline("SanctumAI/Meta-Llama-3.1-8B-Instruct-GGUF")
27
 
28
  def generate_text(model, tokenizer, prompt, max_length=100):
29
  inputs = tokenizer(prompt, return_tensors="pt")
 
40
  if prompt:
41
  col1, col2 = st.columns(2)
42
 
43
+ with col1:
44
+ st.subheader("LLaMA-3.1-8B Output")
45
+ output_8b = pipe8(prompt, max_length)
46
+ st.write(output_8b[0]['generated_text'])
47
 
48
  with col2:
49
  st.subheader("LLaMA-3.1-8B-Instruct Output")
50
+ output_8b_instruct = pipe8instruct(prompt, max_length)
51
+ st.write(output_8b_instruct[0]['generated_text'])
52
  else:
53
  st.warning("Please enter a prompt.")