Update app.py

app.py

@@ -194,7 +194,8 @@ class VideoClassifier:
             except:
                 prompt1 = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
                 messages = [{"role": "user", "content": prompt1}]
-                main_class = self.client.chat_completion(messages, max_tokens=100)
+                stream = self.client.chat_completion(messages, max_tokens=100)
+                main_class = stream.choices[0].message.content.strip()
                 # output = ""
                 # for response in stream:
                 #     output += response['token'].text 
@@ -204,12 +205,15 @@ class VideoClassifier:
                 
             print(main_class)
             print("#######################################################")
-            pattern = r"Main-class =>\s*(.+)"
-            match = re.search(pattern, main_class)
-            if match:
-                main_class = match.group(1).strip()
-            else:
-                main_class = None
+            try:
+                pattern = r"Main-class =>\s*(.+)"
+                match = re.search(pattern, main_class)
+                if match:
+                    main_class = match.group(1).strip()
+                else:
+                    main_class = None
+            except:
+                main_class=main_class
         else:
             prompt_text = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
             response = self.genai_model.generate_content(contents=prompt_text)
@@ -270,18 +274,20 @@ class VideoClassifier:
             except:
                 prompt2 = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
                 messages = [{"role": "user", "content": prompt1}]
-                answer = self.client.chat_completion(messages, max_tokens=100)
+                stream = self.client.chat_completion(messages, max_tokens=100)
+                answer = stream.choices[0].message.content.strip()
                 
             print("Preprocess Answer",answer)
 
-            
-            pattern = r"Sub-class =>\s*(.+)"
-            match = re.search(pattern, answer)
-            if match:
-                sub_class = match.group(1).strip()
-            else:
-                sub_class = None
-            
+            try:
+                pattern = r"Sub-class =>\s*(.+)"
+                match = re.search(pattern, answer)
+                if match:
+                    sub_class = match.group(1).strip()
+                else:
+                    sub_class = None
+            except:
+                subclass=answer
         else:
             prompt_text2 = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
             response = self.genai_model.generate_content(contents=prompt_text2)