from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# We use Flan-T5-Large because it is "Instruction Tuned" 
# (It knows how to answer questions, unlike base BioGPT)
MODEL_NAME = "google/flan-t5-large"

def load_model():
    print(f"Loading {MODEL_NAME}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
    return MedicalModel(model, tokenizer)

class MedicalModel:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def generate(self, text):
        # Prepare the input for the model
        # We add "Answer this medical question:" to help it focus
        input_text = f"Answer this medical question truthfully: {text}"
        
        inputs = self.tokenizer(input_text, return_tensors="pt")

        # Generate response
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_length=200,
                num_beams=5,           # Helps find the best answer
                early_stopping=True
            )
            
        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return [decoded]

def ethics_checker(user_text, ai_response):
    # Simple check for personal data
    risky_terms = ["phone", "ssn", "credit card", "address", "lastname"]
    for k in risky_terms:
        if k in user_text.lower():
            return False, "PII-RISK"
    return True, "OK"