File size: 6,778 Bytes
baa41dd |
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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 |
import gradio as gr
import torch
from transformers import Idefics3ForConditionalGeneration, AutoProcessor
from PIL import Image
# Global variables
model = None
processor = None
device = None
def get_device():
"""Determine best device to use"""
if torch.cuda.is_available():
return 'cuda:0'
else:
return 'cpu'
def load_model():
"""Load SmolVLM model with proper device handling"""
global model, processor, device
try:
print("Loading SmolVLM TRAC Automation Agent...")
device = get_device()
print(f"Using device: {device}")
model_path = r"C:\Users\keith\OneDrive\Desktop\admin.trac.jobs-DATA\LLaMA-Factory_local\smolvlm_final_merged"
# Load processor first
processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
print("✅ Processor loaded")
# Load model with explicit device placement
if device == 'cuda:0':
# GPU loading
model = Idefics3ForConditionalGeneration.from_pretrained(
model_path,
torch_dtype=torch.bfloat16,
device_map={'': 0}, # Force all components to GPU 0
trust_remote_code=True
)
else:
# CPU loading
model = Idefics3ForConditionalGeneration.from_pretrained(
model_path,
torch_dtype=torch.float32, # Use float32 for CPU
device_map='cpu',
trust_remote_code=True
)
print(f"✅ Model loaded on {device}")
return f"✅ Model loaded successfully on {device}! Ready for TRAC automation."
except Exception as e:
error_msg = f"❌ Error loading model: {str(e)}"
print(error_msg)
return error_msg
def analyze_interface(image, task_type, custom_prompt):
"""Analyze TRAC interface with proper device handling"""
global model, processor, device
if model is None:
return "❌ Please load the model first."
if image is None:
return "❌ Please upload a TRAC screenshot."
try:
# Convert image to RGB
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
image = image.convert("RGB")
# Create task-specific prompts
if task_type == "Longlisting":
prompt = """<image>
Analyze this TRAC interface for LONGLISTING candidates. Identify clickable elements, candidate tables, selection controls, and filtering options. Provide automation steps."""
elif task_type == "Shortlisting":
prompt = """<image>
Analyze this TRAC interface for SHORTLISTING candidates. Identify evaluation controls, shortlist buttons, and approval workflows. Provide automation steps."""
elif task_type == "Interview Setup":
prompt = """<image>
Analyze this TRAC interface for INTERVIEW SETUP. Identify scheduling elements, calendar controls, and interviewer assignment. Provide automation steps."""
else: # Custom
if not custom_prompt.strip():
return "❌ Please enter a custom prompt for analysis."
prompt = f"<image>\n{custom_prompt}"
# Process inputs
inputs = processor(text=prompt, images=[image], return_tensors="pt")
# Move ALL tensors to the same device as model
if device == 'cuda:0':
inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in inputs.items()}
# Generate response
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=250,
do_sample=True,
temperature=0.7,
pad_token_id=processor.tokenizer.eos_token_id if hasattr(processor, 'tokenizer') else None
)
# Decode response
response = processor.decode(outputs[0], skip_special_tokens=True)
# Clean up response
if prompt in response:
response = response.replace(prompt, "").strip()
response = response.replace("<image>", "").strip()
if not response:
response = "Model generated empty response. Try a different screenshot or prompt."
return response
except Exception as e:
error_msg = f"❌ Analysis Error: {str(e)}"
print(error_msg)
return error_msg
def create_app():
"""Create Gradio interface"""
with gr.Blocks(title="SmolVLM TRAC Automation") as demo:
gr.Markdown("""
# 🎯 SmolVLM TRAC Automation Agent
**AI Assistant for HR Administrative Tasks**
- 📋 Longlisting candidates
- ⭐ Shortlisting applications
- 📅 Interview setup & scheduling
""")
with gr.Row():
with gr.Column():
# Model loading
load_btn = gr.Button("🚀 Load Model", variant="primary")
status = gr.Textbox(label="Status", value="Model not loaded")
# Image upload
image_input = gr.Image(label="TRAC Screenshot", type="pil")
# Task selection
task_type = gr.Radio(
choices=["Longlisting", "Shortlisting", "Interview Setup", "Custom"],
value="Longlisting",
label="Task Type"
)
# Custom prompt
custom_prompt = gr.Textbox(
label="Custom Prompt",
placeholder="Describe what to analyze...",
lines=3
)
analyze_btn = gr.Button("🔍 Analyze", variant="primary")
with gr.Column():
result = gr.Textbox(
label="Automation Instructions",
lines=15,
show_copy_button=True
)
# Event handlers
load_btn.click(load_model, outputs=status)
analyze_btn.click(
analyze_interface,
inputs=[image_input, task_type, custom_prompt],
outputs=result
)
return demo
if __name__ == "__main__":
print("🌐 Starting SmolVLM TRAC Automation Interface...")
app = create_app()
app.launch(inbrowser=True) |