studfaceval / app5_selectbox /llama2_prompt.py
MENG21's picture
Upload 68 files
e4fe207 verified
raw
history blame contribute delete
761 Bytes
import streamlit as st
model = st.session_state.model
tokenizer = st.session_state.tokenizer
pipeline_generator = st.session_state.pipeline_generator
def llama_prompt(prompt, pipeline=pipeline_generator):
sequences = pipeline(
prompt,
do_sample=True,
top_k=10,
num_return_sequences=1,
eos_token_id=tokenizer.eos_token_id,
max_length=2048, # max length of output, default=4096
return_full_text=False, # to not repeat the question, set to False
# temperature=0.6, # default=0.
)
res = ""
for seq in sequences:
res += str(seq['generated_text'])
return res
# # Example usage:
# prompt = "What is the meaning of life?"
# result = llama_prompt(prompt)
# print(result)