Update app.py

app.py

@@ -2,15 +2,17 @@ import gradio as gr
 from langchain_community.document_loaders import UnstructuredMarkdownLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_core.documents import Document
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpoint
 from langchain_community.vectorstores import FAISS
-from langchain_community.llms import HuggingFaceHub
 from langchain.prompts import ChatPromptTemplate
 from dotenv import load_dotenv
 import os
 from datetime import datetime
 from skyfield.api import load
 import matplotlib.pyplot as plt
+from io import BytesIO
+from PIL import Image
+
 
 # Load environment variables
 load_dotenv()
@@ -121,12 +123,13 @@ def process_query(query_text: str, vectorstore):
         prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
         prompt = prompt_template.format(context=context_text, question=query_text)
         
-        model = HuggingFaceHub(
-            repo_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
+        model = HuggingFaceEndpoint(
+            endpoint_url="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
             task="text2text-generation",
+            # huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),  # Include if token is required
             model_kwargs={"temperature": 0.5, "max_length": 512}
         )
-        response_text = model.predict(prompt)
+        response_text = model.invoke(prompt)
         
         sources = list(set([doc.metadata.get("source", "") for doc, _ in results]))
         return response_text, sources
@@ -297,6 +300,12 @@ def chat_interface(query_text):
             
             # Generate plot
             fig = plot_pladder(PLadder)
+            buf = BytesIO()
+            fig.savefig(buf, format='png')  # Save figure to buffer as PNG
+            buf.seek(0)
+            img = Image.open(buf)  # Convert to PIL image
+            plt.close(fig)  # Close the figure to free memory
+            return text, img
             
             # Compile response text
             text = "Планетарная лестница: " + ", ".join(PLadder_ru) + "\n"