Spaces:
Sleeping
Sleeping
File size: 2,279 Bytes
e479851 67c83c1 88dbb60 3178105 d0ebe8a 101cd9b 2a3cf84 fe16918 2a3cf84 3178105 e479851 3178105 2a3cf84 88dbb60 101cd9b 88dbb60 101cd9b 2a3cf84 101cd9b 88dbb60 101cd9b 67c83c1 2a3cf84 88dbb60 3178105 2a3cf84 e479851 67c83c1 e479851 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import gradio as gr
from transformers import pipeline
import json
import signal
# Load a small, free, instruction-following model
generator = pipeline("text2text-generation", model="google/flan-t5-large")
# Timeout handling
class TimeoutException(Exception):
pass
def timeout_handler(signum, frame):
raise TimeoutException("Processing took too long. Try a simpler input.")
signal.signal(signal.SIGALRM, timeout_handler)
def generate_test_cases(user_story):
try:
signal.alarm(180) # Set a 3-minute timeout
# Structured prompt for better response
prompt = (
f"Generate structured test cases for the following user story:\n"
f"User Story: {user_story}\n"
f"Provide output in a clear, structured way with a test case title, steps, and expected result."
)
output = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]
# Simple manual post-processing to force JSON format
test_cases = []
cases = output.split("\n\n") # Split into test cases
for i, case in enumerate(cases, start=1):
lines = case.split("\n")
if len(lines) >= 3:
title = lines[0].strip()
steps = [line.strip() for line in lines[1:-1]]
expected_result = lines[-1].strip()
test_cases.append({
"id": i,
"title": title,
"steps": steps,
"expected_result": expected_result
})
if not test_cases:
return "Error: Model did not return structured test cases. Try again."
formatted_output = json.dumps({"test_cases": test_cases}, indent=4)
signal.alarm(0) # Disable timeout if successful
return formatted_output
except TimeoutException:
return "Processing timed out. Please try again with a simpler input."
# Gradio UI
iface = gr.Interface(
fn=generate_test_cases,
inputs=gr.Textbox(lines=5, placeholder="Enter your user story here..."),
outputs="text",
title="AI Test Case Generator",
description="Enter a user story and get structured test cases in JSON format.",
)
if __name__ == "__main__":
iface.launch()
|