Update app.py

app.py

@@ -1,56 +1,62 @@
 import streamlit as st
+import requests
 import os
-from openai import OpenAI
 
-# Set the API key
-API_KEY = os.getenv("NEBIUS_API_KEY")  # Use the secret key set in Hugging Face or locally
-if not API_KEY:
-    st.error("API Key is not set. Please set the NEBIUS_API_KEY environment variable.")
-else:
-    # Initialize the OpenAI client
-    client = OpenAI(
-        base_url="https://api.studio.nebius.ai/v1/chat/completions",
-        api_key=API_KEY
-    )
+# API key from environment variable
+API_KEY = os.environ.get("NEBIUS_API_KEY")
 
-    # Streamlit app title
-    st.title("AI Title Generator")
+if not API_KEY:
+    st.error("API key not found. Please set the `NEBIUS_API_KEY` environment variable.")
 
-    # Text Area Input Bar
-    user_input = st.text_area(
-        label="Enter a description for generating titles:",
-        placeholder="e.g., Man who went to jail for no reason"
-    )
+# Function to call Nebius API
+def generate_response(prompt, api_key):
+    api_url = "https://api.studio.nebius.ai/v1/chat/completions"
+    headers = {"Authorization": f"Bearer {api_key}"}
+    payload = {
+        "model": "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
+        "messages": [
+            {"role": "system", "content": "You are an AI that generates short, creative titles based on the user's input."},
+            {"role": "user", "content": prompt}
+        ],
+        "temperature": 0.6,
+        "max_tokens": 512,
+        "top_p": 0.9,
+        "top_k": 50
+    }
+    
+    # API request
+    response = requests.post(api_url, headers=headers, json=payload)
+    
+    if response.status_code == 200:
+        return response.json()
+    else:
+        st.error(f"Error: {response.status_code}, {response.text}")
+        return None
 
-    # Generate Button
-    if st.button("Generate Titles"):
-        if user_input.strip():
-            try:
-                # API call to OpenAI
-                completion = client.chat.completions.create(
-                    model="nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
-                    messages=[
-                        {"role": "system", "content": "Your task is to generate 3 very short titles based on the user input."},
-                        {"role": "user", "content": user_input}
-                    ],
-                    temperature=0.6,
-                    max_tokens=512,  # Corrected argument name
-                    top_p=0.9,
-                    top_k=50
-                )
+# Streamlit App
+st.title("AI Title Generator")
 
-                # Extracting the generated titles
-                response_content = completion.choices[0].message["content"]
+# Input bar for user prompt
+user_input = st.text_area(
+    label="Enter a description for generating titles:",
+    placeholder="e.g., Man who went to jail for no reason"
+)
 
-                # Output Text Area
+# Generate button
+if st.button("Generate Titles"):
+    if user_input.strip():
+        result = generate_response(user_input, API_KEY)
+        if result:
+            try:
+                # Extracting generated titles
+                assistant_message = result["choices"][0]["message"]["content"]
                 st.text_area(
                     label="Generated Titles:",
-                    value=response_content,
+                    value=assistant_message,
                     height=200,
                     disabled=True
                 )
-
-            except Exception as e:
-                st.error(f"An error occurred: {e}")
-        else:
-            st.warning("Please provide input before clicking Generate.")
+            except KeyError as e:
+                st.error(f"Unexpected response format: {e}")
+    else:
+        st.warning("Please provide input before clicking Generate.")