import streamlit as st
import requests
import os

# Get Together API key from environment variable
TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')

# Streamlit UI
st.title("Zero Shot, Single-Shot, Few Shots Prompting Showcase")

# User input for prompt and examples
prompt_and_examples = st.text_area("Enter your prompt and examples:", """\
  Label the sentences as either "positive", "negative", "mixed", or "neutral":

  Sentence: I can say that there isn't anything I would change.
  Label: positive

  Sentence: I'm not sure about this.
  Sentence: neutral

  Sentence: I liked some parts but I didn't like other parts.
  Label: mixed

  Sentence: I think the background image could have been better.
  Label: negative

  Sentence: I really like it.
  Label:""")

# Together API endpoint
endpoint = 'https://api.together.xyz/inference'

# Make a request to Together API
res = requests.post(endpoint, json={
    "model": 'togethercomputer/RedPajama-INCITE-7B-Base',
    "prompt": prompt_and_examples,
    "top_p": 1,
    "top_k": 40,
    "temperature": 0.8,
    "max_tokens": 1,
    "repetition_penalty": 1,
}, headers={
    "Authorization": f"Bearer {TOGETHER_API_KEY}",
    "User-Agent": "<YOUR_APP_NAME>"
})

# Display the output
st.text("Output:")
st.text(res.json()['output']['choices'][0]['text'])