Update gen_ai.py

gen_ai.py

@@ -19,14 +19,14 @@ label_dict = {
 
 
 class traditional_model:
-    def __init__(self, query):
+    def __init__(self):
         self.model =  AutoModelForSequenceClassification.from_pretrained(model_path)
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
 
-    def predict(self):
+    def predict(self,query):
         self.model.to(device)
         self.model.eval()  
-        inputs = self.tokenizer(self.query, return_tensors="pt", truncation=True, padding=True).to(device)  # Move input to device
+        inputs = self.tokenizer(query, return_tensors="pt", truncation=True, padding=True).to(device)  # Move input to device
         with torch.no_grad():
             outputs = self.model(**inputs)
         predicted_class = torch.argmax(outputs.logits, dim=1).item()