Upload 3 files

llm.py

@@ -0,0 +1,78 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
+import gradio as gr
+import time
+
+model = AutoModelForCausalLM.from_pretrained(
+    "MILVLG/imp-v1-3b",
+    torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+tokenizer = AutoTokenizer.from_pretrained("MILVLG/imp-v1-3b", trust_remote_code=True)
+
+
+def response(USER_DATA, TOKEN) -> str:
+    print(USER_DATA)
+    MESSAGE = USER_DATA["text"]
+    NUM_FILES = len(USER_DATA["files"])
+    FILES = USER_DATA["files"]
+
+    SYSTEM_PROMPT = f"""
+    A chat between a curious user and an artificial intelligence assistant. The assistant generates helpful, and detailed testcases for software/website testing.
+    You are tasked with generating detailed, step-by-step test cases for software functionality based on uploaded images. The user will provide one or more images of a software or website interface. For each image, generate a separate set of test cases following the format below:
+
+    Description: Provide a brief explanation of the functionality being tested, as inferred from the image.
+
+    Pre-conditions: Identify any setup requirements, dependencies, or conditions that must be met before testing can begin (e.g., user logged in, specific data pre-populated, etc.).
+
+    Testing Steps: Outline a clear, numbered sequence of actions that a user would take to test the functionality in the image.
+
+    Expected Result: Specify the expected outcome if the functionality is working as intended.
+
+    Ensure that:
+
+    Test cases are created independently for each image.
+    The functionality from each image is fully covered in its own set of test cases.
+    Any assumptions you make are clearly stated.
+    The focus is on usability, navigation, and feature correctness as demonstrated in the UI of the uploaded images.
+
+    USER: <image>\n{MESSAGE}
+    ASSISTANT:
+    """
+
+    RES = generate_answer(FILES, SYSTEM_PROMPT)
+
+    response = f"{RES}."
+    return response
+    for i in range(len(response)):
+        time.sleep(0.025)
+        yield response[: i + 1]
+
+
+def generate_answer(IMAGES: list, SYSTEM_PROMPT) -> str:
+    print(len(IMAGES))
+
+    INPUT_IDS = tokenizer(SYSTEM_PROMPT, return_tensors="pt").input_ids
+
+    RESULT = ""
+    for EACH_IMG in IMAGES:
+        image_path = EACH_IMG["path"]
+        image = Image.open(image_path)
+        image_tensor = model.image_preprocess(image)
+
+        output_ids = model.generate(
+            inputs=INPUT_IDS,
+            max_new_tokens=500,
+            images=image_tensor,
+            use_cache=False,
+        )[0]
+        CUR_RESULT = tokenizer.decode(
+            output_ids[INPUT_IDS.shape[1] :], skip_special_tokens=True
+        ).strip()
+
+        RESULT = f"{RESULT} /n/n {CUR_RESULT}"
+
+    return RESULT