uzmi-ui / src /streamlit_app.py
Rajan Singh
Done
f09166e
raw
history blame contribute delete
691 Bytes
import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
st.title("Uzmi GPT - Romantic Quote Generator")
@st.cache_resource
def load_model():
tokenizer = AutoTokenizer.from_pretrained("rajan3208/uzmi-gpt")
model = AutoModelForCausalLM.from_pretrained("rajan3208/uzmi-gpt")
return tokenizer, model
tokenizer, model = load_model()
prompt = st.text_area("Enter a prompt", "A romantic quote about forever")
if st.button("Generate"):
inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(**inputs, max_new_tokens=50)
generated = tokenizer.decode(output[0], skip_special_tokens=True)
st.success(generated)