update.py

app.py

@@ -1,5 +1,5 @@
 import logging
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
 
 # Set up logging
@@ -9,16 +9,20 @@ logging.basicConfig(
     format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
-# Load the generative AI model
-logging.info("Loading the Hugging Face model...")
+# Select the model you want to use (LLaMA, GPT, or CodeT5)
+MODEL_NAME = "meta-llama/LLaMA-2-7b-hf"
+
+# Load the model and tokenizer
+logging.info(f"Loading the model: {MODEL_NAME}...")
 try:
-    model = pipeline("text2text-generation", model="google/flan-t5-large")  # Replace with your chosen model
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
     logging.info("Model loaded successfully.")
 except Exception as e:
     logging.error(f"Error loading the model: {e}")
     raise
 
-# Function to generate test cases
+# Define a function to generate test cases
 def generate_test_cases(api_info):
     logging.info(f"Generating test cases for API info: {api_info}")
     try:
@@ -26,9 +30,14 @@ def generate_test_cases(api_info):
             f"Generate API test cases for the following API:\n\n{api_info}\n\n"
             f"Test cases should include:\n- Happy path\n- Negative tests\n- Edge cases"
         )
-        result = model(prompt, max_length=512, num_return_sequences=1)
-        logging.info(f"Test cases generated successfully.")
-        return result[0]['generated_text']
+        # Tokenize the input prompt
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+        # Generate output from the model
+        outputs = model.generate(inputs["input_ids"], max_length=512, num_return_sequences=1, do_sample=True)
+        # Decode the generated text
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logging.info("Test cases generated successfully.")
+        return generated_text
     except Exception as e:
         logging.error(f"Error generating test cases: {e}")
         return "An error occurred while generating test cases."
@@ -55,7 +64,7 @@ interface = gr.Interface(
         gr.Textbox(label="Payload (JSON format)"),
     ],
     outputs="text",
-    title="API Test Case Generator"
+    title=f"API Test Case Generator ({MODEL_NAME})"
 )
 
 # Launch Gradio app