modify app

app.py

@@ -6,6 +6,8 @@ import pandas_ta as ta
 from datetime import datetime, timedelta
 import requests
 import json
+import os
+from transformers import pipeline
 
 
 def get_option_stats(ticker):
@@ -168,10 +170,17 @@ def prompt_generator(ticker, option_stats, last_day_summary):
     return output_str
 
 
+model_name = "Rahulholla/mistral-stock-model"
+model = pipeline("text-generation", model=model_name, device_map="auto")
+
+
 def send_prompt_to_api(prompt):
-    url = "https://nu3q3ow9uc.execute-api.us-east-1.amazonaws.com/test/"
-    response = requests.post(url, json={"inputs": prompt})
-    return response
+    try:
+        generated_text = model(prompt, max_length=625)[0]["generated_text"]
+        return {"status_code": 200, "text": generated_text}
+
+    except Exception as e:
+        return {"status_code": 500, "error_message": str(e)}
 
 
 st.title('Stock Tracker, Analysis and Results (S.T.A.R)')