Update app.py

app.py

@@ -192,14 +192,14 @@ class VideoClassifier:
                 chain1 = LLMChain(llm=self.llm, prompt=prompt1)
                 main_class  = chain1.predict(main_categories=main_categories, transcript=transcript, captions=captions)
             except:
-                prompt_text = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
-                stream = self.client.text_generation(prompt_text, **generate_kwargs, stream=True, details=True)
-                output = ""
-                for response in stream:
-                    output += response['token'].text 
-                    print("Streaming output:", output) 
+                messages = [{"role": "user", "content": prompt1}]
+                main_class = self.client.chat_completion(messages, max_tokens=100)
+                # output = ""
+                # for response in stream:
+                #     output += response['token'].text 
+                #     print("Streaming output:", output) 
     
-                main_class = output.strip()
+                # main_class = output.strip()
                 
             print(main_class)
             print("#######################################################")
@@ -267,13 +267,9 @@ class VideoClassifier:
                 chain2 = LLMChain(llm=self.llm, prompt=prompt2)
                 answer = chain2.predict(sub_categories=sub_categories, transcript=transcript, captions=captions,main_class=main_class)
             except:
-                stream = self.client.text_generation(prompt2, **generate_kwargs, stream=True, details=True)
-                output = ""
-                for response in stream:
-                    output += response['token'].text 
-                    print("Streaming output:", output) 
-    
-                main_class = output.strip()
+                messages = [{"role": "user", "content": prompt1}]
+                answer = self.client.chat_completion(messages, max_tokens=100)
+                
             print("Preprocess Answer",answer)
 
             
@@ -308,7 +304,7 @@ class VideoClassifier:
         self.model_name = model_name
         if self.model_name=='mistral':
             print("Setting up Mistral model for Class Selection")
-            self.setup_mistral_model()
+            self.setup_mistral_space_model()
         else :
             print("Setting up Gemini model for Class Selection")
             self.setup_gemini_model()