Update app.py

app.py

@@ -20,6 +20,8 @@ import argparse
 import shutil
 from PIL import Image
 import google.generativeai as genai
+from huggingface_hub import InferenceClient
+
 
 class VideoClassifier:
     def __init__(self, no_of_frames, mode='interface',model='gemini'):
@@ -50,6 +52,16 @@ class VideoClassifier:
         self.whisper_model = whisper.load_model("base")
         self.img_cap = Caption()
 
+    def setup_mistral_space_model(self):
+    if not self.hf_key:
+        raise ValueError("Hugging Face API key is not set or invalid.")
+
+    self.client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2", api_token=self.hf_key)
+
+    self.whisper_model = whisper.load_model("base")
+    self.img_cap = Caption()
+
+    
     def setup_mistral_model(self):
         self.model_id = "mistralai/Mistral-7B-Instruct-v0.2"
         self.device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
@@ -106,7 +118,17 @@ class VideoClassifier:
         print("TRANSCRIPT",transcript)
         return transcript
        
-
+    def generate_text(self, inputs, parameters=None):
+        if parameters is None:
+            parameters = {
+                "temperature": 0.7,
+                "max_new_tokens": 50,
+                "top_p": 0.9,
+                "repetition_penalty": 1.2
+            }
+            
+        return self.client(inputs, parameters)
+    
     def classify_video(self,video_input):
     
         transcript=self.audio_extraction(video_input)
@@ -157,8 +179,13 @@ class VideoClassifier:
         # print(self.model)
         # print(f"Current model in use: {self.model}")
         if self.model_name=='mistral':
-            chain1 = LLMChain(llm=self.llm, prompt=prompt1)
-            main_class  = chain1.predict(main_categories=main_categories, transcript=transcript, captions=captions)
+            try:
+                chain1 = LLMChain(llm=self.llm, prompt=prompt1)
+                main_class  = chain1.predict(main_categories=main_categories, transcript=transcript, captions=captions)
+            except:
+                response = self.client(text=prompt1)
+                main_class = response['generated_text']
+                
             print(main_class)
             print("#######################################################")
             pattern = r"Main-class =>\s*(.+)"
@@ -221,8 +248,12 @@ class VideoClassifier:
         prompt2 = PromptTemplate(template=template2, input_variables=['sub_categories', 'transcript', 'captions','main_class'])
         
         if self.model_name=='mistral':
-            chain2 = LLMChain(llm=self.llm, prompt=prompt2)
-            answer = chain2.predict(sub_categories=sub_categories, transcript=transcript, captions=captions,main_class=main_class)
+            try:
+                chain2 = LLMChain(llm=self.llm, prompt=prompt2)
+                answer = chain2.predict(sub_categories=sub_categories, transcript=transcript, captions=captions,main_class=main_class)
+            except:
+                response = self.client(text=prompt2)
+                answer = response['generated_text']
             print("Preprocess Answer",answer)
 
             
@@ -257,7 +288,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()