File size: 1,024 Bytes
0baf828 |
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 |
import streamlit as st
import os
import litellm
from litellm import completion
# set env variable
os.environ["TOGETHERAI_API_KEY"] = "f722a9f6e3afd6b9999e6aee02aeac9e751ea3a67b124c3667ab50c85c7fa99e"
litellm.register_prompt_template(
model="teknium/OpenHermes-2p5-Mistral-7B",
roles={"system":"<|im_start|>system", "assistant":"<|im_start|>assistant", "user":"<|im_start|>user"}, # tell LiteLLM how you want to map the openai messages to this model
pre_message_sep= "\n",
post_message_sep= "\n"
)
def generate_response(user_input):
messages=[{"role":"user", "content": user_input}]
response = completion(model="together_ai/teknium/OpenHermes-2p5-Mistral-7B", messages=messages)
return response
def main():
st.title("Poem Generator")
user_input = st.text_input("Enter a topic for your poem:")
if st.button("Generate Poem"):
response = generate_response(user_input)
st.write("Here is your poem:")
st.write(response)
if __name__ == "__main__":
main() |