import torchfrom transformers import AutoModelForCausalLM, AutoTokenizerfrom peft import PeftModelimport json model_repo = "center-of-excellence/extract-prompt-quality-criteria"base_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Auto-detect best devicedevice = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")print(f"Using device: {device}") # Load tokenizertokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # Load model — merge LoRA weights for faster inferencebase_model = AutoModelForCausalLM.from_pretrained( base_model_name, dtype=torch.float16 if device in ("cuda", "mps") else torch.float32, trust_remote_code=True, low_cpu_mem_usage=True)model = PeftModel.from_pretrained(base_model, model_repo)model = model.merge_and_unload() # merge for faster inferencemodel = model.to(device)model.eval() # Analyze a promptprompt = "Classify text into categories" instruction = """You are a prompt quality analyzer. Analyze the given prompt and extract quality scores for each criterion. Criteria (score 1-10):1. clarity_and_specificity: Are instructions clear and specific?2. context_sufficiency: Is sufficient background/context provided?3. examples_provided: Are concrete examples included?4. output_format_specification: Is the expected output format clearly defined?5. edge_case_handling: Are edge cases and exceptions addressed?6. tone_and_style_guidance: Is communication style/tone specified?7. constraint_definition: Are limits and boundaries clearly set?8. relevance_of_examples: Do examples match the domain/task? Respond with ONLY a JSON object containing the scores.""" formatted = f"""<|system|>{instruction}<|end|><|user|>Analyze this prompt: {prompt}<|end|><|assistant|>""" inputs = tokenizer(formatted, return_tensors="pt", truncation=True, max_length=512).to(device) with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=300, temperature=0.3, do_sample=True, top_p=0.9, repetition_penalty=1.1, pad_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id ) input_length = inputs['input_ids'].shape[1]response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True) # Parse JSONstart = response.find('{')end = response.rfind('}') + 1result = json.loads(response[start:end])print(json.dumps(result, indent=2))