File size: 1,564 Bytes
75c0510 99d8abc 75c0510 99d8abc 359eb40 99d8abc 359eb40 99d8abc 359eb40 99d8abc 359eb40 99d8abc 359eb40 99d8abc 75c0510 99d8abc 75c0510 99d8abc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import streamlit as st
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
def load_model(model_name):
try:
st.title('Trying to load tokenizer')
tokenizer = AutoTokenizer.from_pretrained(model_name)
st.title('Tokenizer loaded, trying to load model')
model = AutoModelForCausalLM.from_pretrained(model_name)
st.title('Model loaded, initializing pipeline')
generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
st.title('Pipeline ready')
return generator
except Exception as e:
st.error(f"Failed to load model {model_name}: {str(e)}")
return None
# Using distilgpt2 as the model
model_name = "distilgpt2"
generator = load_model(model_name)
# User prompt input
if generator: # Proceed only if the model is successfully loaded
user_prompt = st.text_area("Enter your prompt here:")
# Button to generate text
if st.button('Generate'):
if user_prompt:
# Generate response
try:
response = generator(user_prompt, max_length=50, clean_up_tokenization_spaces=True)
# Display the generated text
st.text_area("Response:", value=response[0]['generated_text'], height=250, disabled=True)
except Exception as e:
st.error(f"Error generating response: {str(e)}")
else:
st.warning("Please enter a prompt.")
else:
st.error("Model could not be loaded. Please ensure the model name is correct and try again.")
|