Update app.py

app.py

@@ -1,3 +1,5 @@
+import os
+
 import gradio as gr
 import torch
 import numpy as np
@@ -10,6 +12,9 @@ from efficientat.helpers.utils import NAME_TO_WIDTH, labels
 from torch import autocast
 from contextlib import nullcontext
 
+from langchain import OpenAI, ConversationChain, LLMChain, PromptTemplate
+from langchain.chains.conversation.memory import ConversationalBufferWindowMemory
+
 MODEL_NAME = "mn40_as"
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
@@ -44,41 +49,48 @@ def audio_tag(
     sorted_indexes = np.argsort(preds)[::-1]
     output = {}
     # Print audio tagging top probabilities
-    for k in range(10):
-        output[sorted_indexes[k]] = labels[sorted_indexes[k]]
+
+    label = labels[sorted_indexes[0]]
+    return formatted_message(label)
+
+
+cached_audio_class = None 
+template = None
+prompt = None
+chain = None
+
+def formatted_message(audio_class):
+    if cached_audio_class != audio_class:
+        cached_audio_class = audio_class
     
-    return "\n".join(output.values())
-
-def formatted_message(audio_length, audio_class, userText):
-    prefix = '''You are going to act as a magical tool that allows for humans to communicate with non-human entities like 
-    rocks, crackling fire, trees, animals, and the wind. In order to do this, we're going to provide you a data string which 
-    represents the audio input, the source of the audio, and the human's text input for the conversation. 
-    The goal is for you to embody the source of the audio, and use the length and variance in the signal data to produce 
-    plausible responses to the humans input. Remember to embody the the source data. When we start the conversation, 
-    you should generate a "personality profile" for the source and utilize that personality profile in your responses. 
-    Let's begin:'''
-    suffix = f'''Source: {audio_class}
-    Length of Audio in Seconds: {audio_length} 
-    Human Input: {userText}
-    {audio_class} Response:'''
-
-    template = prefix + suffix
-
-    response = call_api(template)
-
-    return response
-
-
-def call_api(message):
-    """
-    response = requests.get(f'{api}?q={message}')
-    if response.status_code == 200:
-        
-        return str(response.text).split('\n', 2)[2]
-    else:
-        return Sorry, I'm quite busy right now, but please try again later :)
-    """
-    return message
+        prefix = '''You are going to act as a magical tool that allows for humans to communicate with non-human entities like 
+        rocks, crackling fire, trees, animals, and the wind. In order to do this, we're going to provide you a data string which 
+        represents the audio input, the source of the audio, and the human's text input for the conversation. 
+        The goal is for you to embody the source of the audio, and use the length and variance in the signal data to produce 
+        plausible responses to the humans input. Remember to embody the the source data. When we start the conversation, 
+        you should generate a "personality profile" for the source and utilize that personality profile in your responses. 
+        Let's begin:'''
+        suffix = f'''Source: {audio_class}
+        Length of Audio in Seconds: {audio_length} 
+        Human Input: {userText}
+        {audio_class} Response:'''
+        template = prefix + suffix
+
+        prompt = PromptTemplate(
+            input_variables=["history", "human_input"], 
+            template=template
+        )
+
+        chatgpt_chain = LLMChain(
+            llm=OpenAI(temperature=.5, openai_api_key=session_token), 
+            prompt=prompt, 
+            verbose=True, 
+            memory=ConversationalBufferWindowMemory(k=2),
+        )
+
+    output = chatgpt_chain.predict(human_input=message)
+
+    return output
 
 
 demo = gr.Interface(