Update app.py

app.py

@@ -164,6 +164,15 @@ class VideoClassifier:
         'News and Politics', 'Personal Finance', 'Pets', 'Pop Culture','Real Estate', 'Religion & Spirituality', 'Science', 'Shopping', 'Sports', 
         'Style & Fashion','Technology & Computing', 'Television', 'Travel', 'Video Gaming']
 
+        generate_kwargs = {
+            "temperature": 0.9,
+            "max_new_tokens": 256,
+            "top_p": 0.95,
+            "repetition_penalty": 1.0,
+            "do_sample": True,
+            "seed": 42,
+            "return_full_text": False
+        }
 
         template1 = '''Given below are the different type of main video classes  
         {main_categories}
@@ -183,8 +192,13 @@ class VideoClassifier:
                 chain1 = LLMChain(llm=self.llm, prompt=prompt1)
                 main_class  = chain1.predict(main_categories=main_categories, transcript=transcript, captions=captions)
             except:
-                response = self.client.call_api(task="text-generation", inputs={"inputs": prompt1})
-                main_class = response[0]['generated_text']
+                stream = self.client.text_generation(prompt1, **generate_kwargs, stream=True, details=True)
+                output = ""
+                for response in stream:
+                    output += response['token'].text 
+                    print("Streaming output:", output) 
+    
+                main_class = output.strip()
                 
             print(main_class)
             print("#######################################################")
@@ -252,10 +266,13 @@ 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:
-                # response = self.client(text=prompt2)
-                # answer = response['generated_text']
-                response = self.client.call_api(task="text-generation", inputs={"inputs": prompt2})
-                main_class = response[0]['generated_text']
+                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()
             print("Preprocess Answer",answer)