|
import streamlit as st |
|
import os |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
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') |
|
|
|
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() |
|
|
|
prompt = st.text_input("Enter a prompt:", "Once upon a time") |
|
|
|
if st.button("Generate Text"): |
|
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) |
