Update app.py

app.py

@@ -213,39 +213,41 @@ class VideoClassifier:
         Return only the answer chosen from list and nothing else
         Main-class =>  '''
 
-        prompt1 = PromptTemplate(template=template1, input_variables=['main_categories', 'transcript', 'captions'])
-        print("PROMPT 1",prompt1)
+        formatted_prompt = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
+        print("Formatted prompt:", formatted_prompt)
+        
         # print(self.model)
         # print(f"Current model in use: {self.model}")
         if self.model_name=='mistral':
             try:
+                prompt1 = PromptTemplate(template=template1, input_variables=['main_categories', 'transcript', 'captions'])
+                print("PROMPT 1",prompt1)
                 print("Entering mistral chain approach")
                 chain1 = LLMChain(llm=self.llm, prompt=prompt1)
                 main_class  = chain1.predict(main_categories=main_categories, transcript=transcript, captions=captions)
             except:
-                print("Entering mistral template approach")
                 output_queue = queue.Queue()
-                # prompt1 = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
-                # messages = [{"role": "user", "content": prompt1}]
-                # stream = self.client.chat_completion(messages, max_tokens=100)
-                # main_class = stream.choices[0].message.content.strip()
+
+                def get_classification():
+                    messages = [{"role": "user", "content": formatted_prompt}]
+                    try:
+                        stream = self.client.chat_completion(messages, max_tokens=100)
+                        main_class = stream.choices[0].message.content.strip()
+                    except Exception as ex:
+                        main_class = f"Error: {str(ex)}"
+                    output_queue.put(main_class)
+    
                 classification_thread = threading.Thread(target=get_classification)
                 classification_thread.start()
                 classification_thread.join(timeout=30)
+    
                 if classification_thread.is_alive():
                     print("Classification timeout occurred.")
                     return "Timeout or error during classification."
-
+    
                 # Get result from queue
                 main_class = output_queue.get()
-                print("MAIN CLASS template:", main_class)
-                return main_class
-                # output = ""
-                # for response in stream:
-                #     output += response['token'].text 
-                #     print("Streaming output:", output) 
-    
-                # main_class = output.strip()
+                print("MAIN CLASS:", main_class)
                 
             print(main_class)
             print("#######################################################")