Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import openai
-
 import os
 
 # Securely fetch the OpenAI API key
@@ -45,16 +44,14 @@ KNOWN_MODELS = [
 
 def recommend_ai_model_via_gpt(description):
     messages = [
-        {"role": "user", "content": f"Given the application described as: '{description}', which AI model would be most suitable?"}
+        {"role": "user", "content": description}
     ]
-
-    #try:
     response = openai.ChatCompletion.create(
         model="gpt-4",
         messages=messages
     )
-    recommendation = response['choices'][0]['message']['content'].strip()
-    return recommendation
+    return response['choices'][0]['message']['content'].strip()
+
     #except openai.APIError as e:
     #    return f"Error: {e}"
     #except openai.RateLimitError as e:
@@ -67,14 +64,11 @@ def explain_recommendation(model_name):
     messages = [
         {"role": "user", "content": f"Why would {model_name} be a suitable choice for the application?"}
     ]
-    
-    #try:
     response = openai.ChatCompletion.create(
         model="gpt-4",
         messages=messages
     )
-    explanation = response['choices'][0]['message']['content'].strip()
-    return explanation
+    return response['choices'][0]['message']['content'].strip()
     #except openai.APIError as e:
     #    return f"Error: {e}"
     #except openai.RateLimitError as e:
@@ -86,7 +80,9 @@ def explain_recommendation(model_name):
 st.image("./A8title2.png")
 st.title('Find the best AI stack for your app')
 
-description = st.text_area("Describe your application:", "")
+description = st.text_area("Describe your application:")
+dataset_description = st.text_area("Describe the dataset you want to use for fine-tuning your model:")
+
 recommendation_type = st.radio("What type of recommendation are you looking for?", ["Recommend Open-Source Model", "Recommend API Service"])
 
 if "rec_model_pressed" not in st.session_state:
@@ -99,25 +95,13 @@ if st.button("Recommend AI Model"):
     st.session_state.rec_model_pressed = True
 
 if st.session_state.rec_model_pressed:
-    if description:
-        # Modified query based on recommendation type
-        if recommendation_type == "Recommend Open-Source Model":
-            query = f"Given the application described as: '{description}', which open-source AI model would be most suitable?"
-        else:  # Recommend API Service
-            query = f"Given the application described as: '{description}', which AI service API would be best?"
-
-        recommended_model = recommend_ai_model_via_gpt(query)  # Updated function call
-
-        # Validate recommended model
-        # Commenting out model validation for the example
-        # if recommended_model not in KNOWN_MODELS:
-        #     st.warning("The recommendation is ambiguous. Please refine your description or consult an expert.")
-        # else:
+    if description and dataset_description:
+        combined_query = f"{description} Dataset: {dataset_description}"
+        recommended_model = recommend_ai_model_via_gpt(combined_query)
         st.subheader(f"Recommended: {recommended_model}")
         explanation = explain_recommendation(recommended_model)
         st.write("Reason:", explanation)
         
-        # Collecting rating and feedback through Streamlit
         rating = st.slider("Rate the explanation from 1 (worst) to 5 (best):", 1, 5)
         feedback = st.text_input("Any additional feedback?")
 
@@ -127,9 +111,5 @@ if st.session_state.rec_model_pressed:
         if st.session_state.feedback_submitted:
             st.success("Thank you for your feedback!")
 
-            st.write("Contact team@autumn8.ai or call (857) 600-0180 to learn how we can fine-tune and host this app for you.")
-            
-    else:
-        st.warning("Please provide a description.")
-
-
+else:
+    st.warning("Please provide a description and dataset details.")