|
import streamlit as st |
|
import os |
|
import logging |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
|
logging.basicConfig(level=logging.INFO) |
|
logger = logging.getLogger(__name__) |
|
|
|
st.title("Meta LLaMA Text Generation") |
|
|
|
@st.cache_resource |
|
def load_model(): |
|
model_name = "meta-llama/Meta-Llama-3-8B" |
|
access_token = os.getenv('hf') |
|
|
|
if not access_token: |
|
st.error("Hugging Face access token is not set. Please set the environment variable 'hf'.") |
|
return None, None |
|
|
|
logger.info("Loading tokenizer and model...") |
|
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=access_token) |
|
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=access_token) |
|
return tokenizer, model |
|
|
|
tokenizer, model = load_model() |
|
|
|
if tokenizer is not None and model is not None: |
|
prompt = st.text_input("Enter a prompt:", "Once upon a time") |
|
|
|
if st.button("Generate Text"): |
|
try: |
|
inputs = tokenizer(prompt, return_tensors="pt") |
|
outputs = model.generate(**inputs, max_length=50) |
|
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
st.write(generated_text) |
|
except Exception as e: |
|
st.error(f"An error occurred: {e}") |
|
logger.error(f"An error occurred during text generation: {e}") |
|
else: |
|
st.error("Failed to load the model. Check the logs for more details.") |
|
|
