Inferences on Spaces/Streamlit

Spaces
1

Hey ! I am trying to build a Space with Streamlit where you can run inferences, but my code doesn’t seem to be able to complete the generate() function. The code goes to st.write(f'inst: {instruction}') but not up to t.write('Gen done'). Here is the code:

import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
import warnings
import os
import streamlit as st

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
warnings.filterwarnings("ignore", category=UserWarning, module='transformers.generation.utils')

@st.cache_resource
def load_model():
    base_model = "microsoft/phi-2"
    peft_model_id = "STEM-AI-mtl/phi-2-electrical-engineering"
    config = PeftConfig.from_pretrained(peft_model_id, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(base_model, device_map="cpu", return_dict=True, trust_remote_code=True)
    model = PeftModel.from_pretrained(model, peft_model_id, trust_remote_code=True)
    #model = model.to('cpu')
    return model

@st.cache_resource
def load_tokenizer():
    base_model = "microsoft/phi-2"
    return AutoTokenizer.from_pretrained(base_model)

model = load_model()
tokenizer = load_tokenizer()

def generate(instruction, model, tokenizer):                              
    inputs = tokenizer(instruction, return_tensors="pt", return_attention_mask=False)
    #inputs = inputs.to('cpu')
    st.write(f'inst: {instruction}')
    outputs = model.generate(
        **inputs, 
        max_length=350,
        do_sample=True, 
        temperature=0.7,
        top_k=50,  
        top_p=0.9,
        repetition_penalty=1,
    )
    st.write('Gen done')
    text = tokenizer.batch_decode(outputs)[0]
    st.write('Decode done')
    return text
    
# Streamlit interface
st.title('Electrical_engineer.AI')

instruction = st.text_input("Enter your instruction:")

if st.button('Submit'):
    if instruction:
        answer = generate(instruction, model, tokenizer)
        st.write(f'Answer: {answer}')
    else:
        st.write("Please enter an instruction.")