Switched to pipeline api

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 import torch
 import os
 
@@ -11,15 +11,19 @@ if not hf_token:
     st.error("Hugging Face token not found. Please add your HF_TOKEN to the Space secrets.")
     st.stop()
 
-# Load models and tokenizers
-@st.cache_resource
-def load_model_and_tokenizer(model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    return model, tokenizer
 
-model_8b, tokenizer_8b = load_model_and_tokenizer("unsloth/Meta-Llama-3.1-8B-bnb-4bit")
-model_8b_instruct, tokenizer_8b_instruct = load_model_and_tokenizer("SanctumAI/Meta-Llama-3.1-8B-Instruct-GGUF")
+@st.cache_resource
+def load_pipeline(model_name):
+    with st.spinner(f'Loading {model_name}... This may take several minutes.'):
+        try:
+            pipe = pipeline("text-generation", model=model_name)
+        except Exception as e:
+            st.error(f"An error occurred: {e}")
+            st.stop()
+    return pipe
+
+pipe8 = load_pipeline("unsloth/Meta-Llama-3.1-8B-bnb-4bit")
+pipe8instruct = load_pipeline("SanctumAI/Meta-Llama-3.1-8B-Instruct-GGUF")
 
 def generate_text(model, tokenizer, prompt, max_length=100):
     inputs = tokenizer(prompt, return_tensors="pt")
@@ -36,14 +40,14 @@ if st.button("Generate"):
     if prompt:
         col1, col2 = st.columns(2)
         
-        # with col1:
-        #     st.subheader("LLaMA-3.1-8B Output")
-        #     output_8b = generate_text(model_8b, tokenizer_8b, prompt, max_length)
-        #     st.write(output_8b)
+        with col1:
+            st.subheader("LLaMA-3.1-8B Output")
+            output_8b = pipe8(prompt, max_length)
+            st.write(output_8b[0]['generated_text'])
         
         with col2:
             st.subheader("LLaMA-3.1-8B-Instruct Output")
-            output_8b_instruct = generate_text(model_8b_instruct, tokenizer_8b_instruct, prompt, max_length)
-            st.write(output_8b_instruct)
+            output_8b_instruct = pipe8instruct(prompt, max_length)
+            st.write(output_8b_instruct[0]['generated_text'])
     else:
         st.warning("Please enter a prompt.")