Update app.py

app.py

@@ -1,7 +1,11 @@
 import gradio as gr
 import time
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-prompt=''' Review the Following text as a human, who is tasked to Extract `Code Snippets`If No code Snippets are there then return NO.This is very important for my career:
+prompt1=''' is there any code Snippets present in the below text? Reply with YES or NO ONLY. This is very important for my career:
+
+Text: " {} "
+'''
+prompt2='''Extract Code Snippets ONLY from the below text .This is very important for my career:
 
 Text: " {} "
 '''
@@ -18,7 +22,14 @@ def chatbot_model(m,history):
     for char, replacement in replacements.items():
         m = m.replace(char, replacement)
     # Encode the prompt and generate response
-    input_ids = tokenizer.encode(prompt.format(m) + tokenizer.eos_token, return_tensors='pt')
+    input_ids = tokenizer.encode(prompt1.format(m) + tokenizer.eos_token, return_tensors='pt')
+    outputs = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+    
+    # Decode the output
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    if (decode_output=='NO'):
+        yield 'NO'
+    input_ids = tokenizer.encode(prompt2.format(m) + tokenizer.eos_token, return_tensors='pt')
     outputs = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
     
     # Decode the output