app.py

import streamlit as st
# from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import LlamaTokenizer, LlamaForCausalLM

# Load your model here
# tokenizer = AutoTokenizer.from_pretrained("qbwmwsap/amber-model-mine")
# model = AutoModelForCausalLM.from_pretrained("qbwmwsap/amber-model-mine")

# Load original amber model
tokenizer = LlamaTokenizer.from_pretrained("LLM360/Amber", revision="ckpt_356")
model = LlamaForCausalLM.from_pretrained("LLM360/Amber", revision="ckpt_356")

st.title('Amber-Model-Mine')

text = st.text_input('Enter some text')

if text:
    # Preprocess your data here
    input_ids = tokenizer(text, return_tensors="pt").input_ids

    # Run the model and get the output
    outputs = model.generate(input_ids)

    # Postprocess your output here if necessary    
    output = tokenizer.decode(outputs[0])

    st.write(output)