| |
| |
| """ |
| ============================================================================== |
| Manara-3B | Inference & Evaluation Script |
| ============================================================================== |
| This script loads the fine-tuned Manara-3B model and runs inference, |
| including the self-reflective critic pass for guardrail compliance. |
| |
| Usage: |
| python inference.py \ |
| --model_path ../../model_weights/manara-3b-final \ |
| --prompt "ما هي متطلبات نسبة كفاية رأس المال وفق تعليمات بنك الكويت المركزي؟" \ |
| --language ar |
| ============================================================================== |
| """ |
|
|
| import sys |
| import argparse |
| import logging |
| from pathlib import Path |
|
|
| import torch |
|
|
| try: |
| from unsloth import FastLanguageModel |
| UNSLOTH_AVAILABLE = True |
| except ImportError: |
| UNSLOTH_AVAILABLE = False |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| |
| sys.path.append(str(Path(__file__).parent.parent / "guardrails")) |
| from self_reflective_loop import ManaraGuardrailEngine |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") |
| logger = logging.getLogger(__name__) |
|
|
| ALPACA_PROMPT_AR = """فيما يلي تعليمات تصف مهمة. اكتب استجابة تكمل الطلب بشكل مناسب. |
| |
| ### التعليمات: |
| {} |
| |
| ### الاستجابة: |
| """ |
|
|
| ALPACA_PROMPT_EN = """Below is an instruction that describes a task. Write a response that appropriately completes the request. |
| |
| ### Instruction: |
| {} |
| |
| ### Response: |
| """ |
|
|
|
|
| class ManaraInferenceEngine: |
| """Handles inference for the Manara-3B model with guardrail integration.""" |
|
|
| def __init__(self, model_path: str, max_seq_length: int = 2048): |
| self.model_path = model_path |
| self.max_seq_length = max_seq_length |
| self._load_model() |
| self.guardrail_engine = ManaraGuardrailEngine() |
|
|
| def _load_model(self): |
| """Loads the fine-tuned model for inference.""" |
| if UNSLOTH_AVAILABLE: |
| self.model, self.tokenizer = FastLanguageModel.from_pretrained( |
| model_name=self.model_path, |
| max_seq_length=self.max_seq_length, |
| dtype=None, |
| load_in_4bit=True, |
| ) |
| FastLanguageModel.for_inference(self.model) |
| else: |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_path) |
| self.model = AutoModelForCausalLM.from_pretrained( |
| self.model_path, |
| device_map="auto", |
| torch_dtype=torch.float16, |
| ) |
| logger.info(f"Model loaded from: {self.model_path}") |
|
|
| def generate( |
| self, |
| prompt: str, |
| language: str = "ar", |
| max_new_tokens: int = 512, |
| temperature: float = 0.7, |
| top_p: float = 0.9, |
| apply_guardrails: bool = True, |
| ) -> dict: |
| """ |
| Generates a response for the given prompt, with optional guardrail checking. |
| |
| Returns a dict containing: |
| - 'response': The generated text |
| - 'guardrail_passed': Whether the response passed all guardrails |
| - 'corrections': Any corrections made by the critic pass |
| """ |
| |
| template = ALPACA_PROMPT_AR if language == "ar" else ALPACA_PROMPT_EN |
| formatted_prompt = template.format(prompt) |
|
|
| |
| inputs = self.tokenizer([formatted_prompt], return_tensors="pt").to("cuda") |
|
|
| |
| with torch.no_grad(): |
| outputs = self.model.generate( |
| **inputs, |
| max_new_tokens=max_new_tokens, |
| temperature=temperature, |
| top_p=top_p, |
| do_sample=True, |
| pad_token_id=self.tokenizer.eos_token_id, |
| ) |
|
|
| |
| generated_ids = outputs[0][inputs["input_ids"].shape[1]:] |
| response = self.tokenizer.decode(generated_ids, skip_special_tokens=True) |
|
|
| result = { |
| "prompt": prompt, |
| "response": response, |
| "language": language, |
| "guardrail_passed": True, |
| "corrections": [], |
| } |
|
|
| |
| if apply_guardrails: |
| guardrail_result = self.guardrail_engine.check(response, language) |
| result["guardrail_passed"] = guardrail_result["passed"] |
| result["corrections"] = guardrail_result["violations"] |
|
|
| if not guardrail_result["passed"]: |
| logger.warning(f"Guardrail violations detected: {guardrail_result['violations']}") |
| |
| |
|
|
| return result |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description="Manara-3B: Inference") |
| parser.add_argument("--model_path", required=True, help="Path to fine-tuned model") |
| parser.add_argument("--prompt", required=True, help="Input prompt") |
| parser.add_argument("--language", default="ar", choices=["ar", "en"]) |
| parser.add_argument("--max_new_tokens", type=int, default=512) |
| parser.add_argument("--no_guardrails", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| if __name__ == "__main__": |
| args = parse_args() |
| engine = ManaraInferenceEngine(args.model_path) |
| result = engine.generate( |
| prompt=args.prompt, |
| language=args.language, |
| max_new_tokens=args.max_new_tokens, |
| apply_guardrails=not args.no_guardrails, |
| ) |
| print("\n" + "=" * 60) |
| print(f"PROMPT: {result['prompt']}") |
| print("=" * 60) |
| print(f"RESPONSE:\n{result['response']}") |
| print("=" * 60) |
| print(f"GUARDRAIL PASSED: {result['guardrail_passed']}") |
| if result["corrections"]: |
| print(f"VIOLATIONS: {result['corrections']}") |
| print("=" * 60) |
|
|
