Spaces:
Running
Running
| import streamlit as st | |
| from openai import OpenAI | |
| import os | |
| # Get API Key from environment variable | |
| api_key = os.getenv("NVIDIA_API_KEY") | |
| # Check if the API key is found | |
| if api_key is None: | |
| st.error("NVIDIA_API_KEY environment variable not found.") | |
| else: | |
| # Initialize the OpenAI client | |
| client = OpenAI( | |
| base_url="https://integrate.api.nvidia.com/v1", | |
| api_key=api_key | |
| ) | |
| def get_response(user_prompt): | |
| try: | |
| # Create and send the chat completion request | |
| completion = client.chat.completions.create( | |
| model="nvidia/llama-3.1-nemotron-70b-instruct", | |
| messages=[{"role": "user", "content": user_prompt}], | |
| temperature=0.5, # Adjust temperature for creativity | |
| top_p=1, | |
| max_tokens=1024, | |
| stream=True # Enable streaming response | |
| ) | |
| # Display the streaming response | |
| response_container = st.empty() | |
| full_response = "" | |
| for chunk in completion: | |
| if chunk.choices[0].delta.content is not None: | |
| full_response += chunk.choices[0].delta.content | |
| response_container.write(full_response) | |
| except Exception as e: | |
| st.error(f"An error occurred: {e}") | |
| def main(): | |
| # Streamlit app | |
| st.title("NVIDIA Nemotron Chatbot") | |
| # User Prompt Input | |
| user_prompt = st.text_area("Enter your prompt:") | |
| # When the user clicks "Submit" | |
| if st.button("Submit"): | |
| if user_prompt: | |
| # Call the OpenAI API | |
| get_response(user_prompt) | |
| else: | |
| st.warning("Please enter a prompt.") | |
| if __name__ == "__main__": | |
| main() | |