Create model..py

model..py

@@ -0,0 +1,41 @@
+from transformers import pipeline
+
+# Function to load the model
+def load_model(model_name):
+    try:
+        # Load the model from Hugging Face or local storage (by name)
+        model = pipeline("text-classification", model=model_name)
+        return model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None
+
+# Function to run inference using the selected model
+def run_inference(user_input, selected_model, prompt=None):
+    model = load_model(selected_model)
+    if model:
+        # If a prompt is provided, prepend it to the input text
+        if prompt:
+            input_text = f"{prompt}\n{user_input}"
+        else:
+            input_text = user_input
+        
+        try:
+            # Run inference and check model output
+            result = model(input_text)
+            
+            # Assuming the output format is a list of dicts with 'label' field
+            return result[0]['label'] if 'label' in result[0] else "Error: No label in output"
+        except Exception as e:
+            return f"Error during inference: {e}"
+    else:
+        return f"Error: Model '{selected_model}' failed to load."
+
+# Example usage
+selected_model = "Canstralian/CySec_Known_Exploit_Analyzer"
+user_input = "Sample exploit description"
+prompt = "Classify the following cybersecurity exploit:"
+
+# Run inference
+result = run_inference(user_input, selected_model, prompt)
+print(f"Inference Result: {result}")