medico2025 / model_functions.py

import error corrected

a9673bf about 1 year ago

7.94 kB

import torch 
from torch.nn import CrossEntropyLoss, MSELoss
import re
from models import disease_model
#!pip install rouge_score

#from rouge_score import rouge_scorer

from nltk.translate.meteor_score import meteor_score


def forward_batch(images, input_ids, attention_mask, answers, question_classes=None,qtype_classifier=None,fusion_module=None,q_types=None,q_types_mapping=None,task_heads=None,device=None,image_encoder=None,question_encoder=None):
    # Image encoding
    img_outputs = image_encoder(pixel_values=images.to(device))
    img_feat = img_outputs.last_hidden_state  # [B, R, 768]

    # Question encoding (DistilBERT for qtype classification)
    task_logits = qtype_classifier(input_ids=input_ids.to(device), 
                                   attention_mask=attention_mask.to(device))  # [B, num_types]

    # Use another encoder for question embeddings (router encoder you already had)
    q_feat = question_encoder(input_ids=input_ids.to(device), 
                              attention_mask=attention_mask.to(device)).pooler_output  # [B, 768]

    # Disease model
    disease_vec = disease_model(images.to(device))  # [B, 23]

    # Fusion
    fused = fusion_module(img_feat, q_feat, disease_vec)

    # Task-specific predictions (list of preds per sample, like before)
    preds = []
    for i, q_class in enumerate(question_classes):#q_class from task)type
        mapped_type = q_types_mapping[q_class[0] if isinstance(q_class, list) else q_class]
        predictor = task_heads[mapped_type]   # ✅ trained head
        pred_out = predictor(fused[i].unsqueeze(0))
        preds.append(pred_out)
        #general_class = q_types_mapping[task_type[0] if isinstance(task_type, list) else task_type]
        #head = TaskPredictor(general_class, hidden=fused.size(-1)).to(device)
        #preds.append(head(fused[i].unsqueeze(0)))

    return preds, answers, task_logits


def forward_batch1(images, input_ids, attention_mask, answers, true_q_classes=None,qtype_classifier=None,fusion_module=None,q_types=None):
    # Disease vector (dummy placeholder: replace with your trained disease model)
    disease_vec = disease_model(images)  # [B, 23]

    # Encode image
    img_outputs = image_encoder(pixel_values=images.to(device))
    img_feat = img_outputs.last_hidden_state  # [B, R, 768]

    # Encode question
    q_feat = question_encoder(input_ids=input_ids.to(device),
                              attention_mask=attention_mask.to(device)).pooler_output  # [B, 768]

    # Predict task type from question
    #print(q_feat.device)
    #print(q_feat.shape)
    #task_logits = qtype_classifier(q_feat)  # [B, 6]
    task_logits = qtype_classifier(input_ids=batch["input_ids"],
                               attention_mask=batch["attention_mask"])
    task_pred = torch.argmax(task_logits, dim=1)  # predicted type index
    
    # Fusion
    fused = fusion_module(img_feat, q_feat, disease_vec)

    # Task-specific predictions
    preds = []
    for i, t_idx in enumerate(task_pred):
        task_type = q_types[t_idx]  # map index to string
        predictor = TaskPredictor(task_type).to(device)
        preds.append(predictor(fused[i].unsqueeze(0)))

    return preds, answers, task_logits
    
    #for i, task_type in enumerate(q_classes):
    #    predictor = TaskPredictor(task_type).to(device)
    #    pred_out = predictor(fused[i].unsqueeze(0))
    #    preds.append(pred_out)
    #return preds, answers


def extract_count(answer_str):
    """
    Try to convert an answer string into a number.
    Returns None if it cannot be parsed.
    """
    try:
        # Direct numeric
        return float(answer_str)
    except ValueError:
        pass

    # Handle words like "one", "two", etc.
    word2num = {
        "zero": 0, "one": 1, "two": 2, "three": 3,
        "four": 4, "five": 5, "six": 6,
        "seven": 7, "eight": 8, "nine": 9, "ten": 10
    }
    tokens = answer_str.lower().split()
    for t in tokens:
        if t in word2num:
            return float(word2num[t])

    # Extract any digits from the string
    numbers = re.findall(r"\d+", answer_str)
    if numbers:
        return float(numbers[0])

    return None  # fallback


def compute_meteor(preds, answers, answer_vocabs, mapped_classes):
    scores = []
    for pred, ans, c in zip(preds, answers, mapped_classes):
        if c not in answer_vocabs:
            continue
        # Get predicted index
        pred_idx = pred.argmax(dim=1).item()
        # Map index back to string
        inv_vocab = {v: k for k, v in answer_vocabs[c].items()}
        pred_str = inv_vocab.get(pred_idx, "")
        # METEOR score between predicted and ground truth answer
        score = meteor_score([ans.split()], pred_str.split())
        scores.append(score)
    return sum(scores) / len(scores) if scores else 0.0


def compute_rouge(preds, answers, answer_vocabs, mapped_classes):
    scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
    scores = []
    for pred, ans, c in zip(preds, answers, mapped_classes):
        if c not in answer_vocabs:
            continue
        pred_idx = pred.argmax(dim=1).item()
        inv_vocab = {v: k for k, v in answer_vocabs[c].items()}
        pred_str = inv_vocab.get(pred_idx, "")
        score = scorer.score(ans, pred_str)["rougeL"].fmeasure
        scores.append(score)
    return sum(scores) / len(scores) if scores else 0.0

def compute_loss(preds, answers, task_logits, true_q_classes, answer_vocabs,q_types_mapping,q_types,task_heads):
    """
    preds: list of model predictions for each sample
    answers: list of strings (descriptive answers)
    task_logits: tensor [batch_size, num_task_types]
    true_q_classes: list of lists (fine-grained classes for each question)
    answer_vocabs: dict mapping {q_type: {answer: index}}
    """

    ce_loss = CrossEntropyLoss()
    mse_loss = MSELoss()

    total_loss = 0

    # 1) Map fine-grained → general classes
    mapped_classes = [
        q_types_mapping[c[0] if isinstance(c, list) else c]
        for c in true_q_classes
    ]

    # 2) Question type classification loss
    true_task_types = torch.tensor(
        [q_types.index(c) for c in mapped_classes],
        device=task_logits.device
    )

    #print("task_logits, true_task_types\t",task_logits, true_task_types)
    
    #print("task_logits, true_task_types\t",task_logits.shape, true_task_types.shape)
    
    task_loss = ce_loss(task_logits, true_task_types)
    total_loss += task_loss

    # 3) Answer prediction loss (per sample)
    for pred, ans, c in zip(preds, answers, mapped_classes):
        predictor = task_heads[c]   # ✅ trained head
        if c == "count":
            # For count, answer must be numeric
            try:
                ans_val = float(ans)
                ans_val = torch.tensor([ans_val], device=pred.device)
                total_loss += mse_loss(pred.squeeze(), ans_val)
            except ValueError:
                print(f"[Warning] Skipping non-numeric count answer: {ans}")
                continue

        else:
            # For categorical tasks (yesno, single, multi, etc.)
            if ans not in answer_vocabs.get(c, {}):
                print(f"[Warning] Skipping unseen or descriptive answer {ans} for task {c}")
                continue

            ans_idx = answer_vocabs[c][ans]

            if ans_idx >= pred.size(1):
                print(f"[Warning] Skipping answer {ans} for task {c}: "
                      f"index {ans_idx} >= pred.size(1)")
                continue

            ans_tensor = torch.tensor([ans_idx], device=pred.device)
            total_loss += ce_loss(pred, ans_tensor)

    meteor = compute_meteor(preds, answers, answer_vocabs, mapped_classes)
    print(f"Validation METEOR: {meteor:.4f}")
    #rouge = compute_rouge(preds, answers, answer_vocabs, mapped_classes)
    #print(f"Validation ROUGE-L: {rouge:.4f}")

    
    return total_loss / len(preds)

	import torch
	from torch.nn import CrossEntropyLoss, MSELoss
	import re
	from models import disease_model
	#!pip install rouge_score

	#from rouge_score import rouge_scorer

	from nltk.translate.meteor_score import meteor_score


	def forward_batch(images, input_ids, attention_mask, answers, question_classes=None,qtype_classifier=None,fusion_module=None,q_types=None,q_types_mapping=None,task_heads=None,device=None,image_encoder=None,question_encoder=None):
	# Image encoding
	img_outputs = image_encoder(pixel_values=images.to(device))
	img_feat = img_outputs.last_hidden_state # [B, R, 768]

	# Question encoding (DistilBERT for qtype classification)
	task_logits = qtype_classifier(input_ids=input_ids.to(device),
	attention_mask=attention_mask.to(device)) # [B, num_types]

	# Use another encoder for question embeddings (router encoder you already had)
	q_feat = question_encoder(input_ids=input_ids.to(device),
	attention_mask=attention_mask.to(device)).pooler_output # [B, 768]

	# Disease model
	disease_vec = disease_model(images.to(device)) # [B, 23]

	# Fusion
	fused = fusion_module(img_feat, q_feat, disease_vec)

	# Task-specific predictions (list of preds per sample, like before)
	preds = []
	for i, q_class in enumerate(question_classes):#q_class from task)type
	mapped_type = q_types_mapping[q_class[0] if isinstance(q_class, list) else q_class]
	predictor = task_heads[mapped_type] # ✅ trained head
	pred_out = predictor(fused[i].unsqueeze(0))
	preds.append(pred_out)
	#general_class = q_types_mapping[task_type[0] if isinstance(task_type, list) else task_type]
	#head = TaskPredictor(general_class, hidden=fused.size(-1)).to(device)
	#preds.append(head(fused[i].unsqueeze(0)))

	return preds, answers, task_logits


	def forward_batch1(images, input_ids, attention_mask, answers, true_q_classes=None,qtype_classifier=None,fusion_module=None,q_types=None):
	# Disease vector (dummy placeholder: replace with your trained disease model)
	disease_vec = disease_model(images) # [B, 23]

	# Encode image
	img_outputs = image_encoder(pixel_values=images.to(device))
	img_feat = img_outputs.last_hidden_state # [B, R, 768]

	# Encode question
	q_feat = question_encoder(input_ids=input_ids.to(device),
	attention_mask=attention_mask.to(device)).pooler_output # [B, 768]

	# Predict task type from question
	#print(q_feat.device)
	#print(q_feat.shape)
	#task_logits = qtype_classifier(q_feat) # [B, 6]
	task_logits = qtype_classifier(input_ids=batch["input_ids"],
	attention_mask=batch["attention_mask"])
	task_pred = torch.argmax(task_logits, dim=1) # predicted type index

	# Fusion
	fused = fusion_module(img_feat, q_feat, disease_vec)

	# Task-specific predictions
	preds = []
	for i, t_idx in enumerate(task_pred):
	task_type = q_types[t_idx] # map index to string
	predictor = TaskPredictor(task_type).to(device)
	preds.append(predictor(fused[i].unsqueeze(0)))

	return preds, answers, task_logits

	#for i, task_type in enumerate(q_classes):
	# predictor = TaskPredictor(task_type).to(device)
	# pred_out = predictor(fused[i].unsqueeze(0))
	# preds.append(pred_out)
	#return preds, answers


	def extract_count(answer_str):
	"""
	Try to convert an answer string into a number.
	Returns None if it cannot be parsed.
	"""
	try:
	# Direct numeric
	return float(answer_str)
	except ValueError:
	pass

	# Handle words like "one", "two", etc.
	word2num = {
	"zero": 0, "one": 1, "two": 2, "three": 3,
	"four": 4, "five": 5, "six": 6,
	"seven": 7, "eight": 8, "nine": 9, "ten": 10
	}
	tokens = answer_str.lower().split()
	for t in tokens:
	if t in word2num:
	return float(word2num[t])

	# Extract any digits from the string
	numbers = re.findall(r"\d+", answer_str)
	if numbers:
	return float(numbers[0])

	return None # fallback


	def compute_meteor(preds, answers, answer_vocabs, mapped_classes):
	scores = []
	for pred, ans, c in zip(preds, answers, mapped_classes):
	if c not in answer_vocabs:
	continue
	# Get predicted index
	pred_idx = pred.argmax(dim=1).item()
	# Map index back to string
	inv_vocab = {v: k for k, v in answer_vocabs[c].items()}
	pred_str = inv_vocab.get(pred_idx, "")
	# METEOR score between predicted and ground truth answer
	score = meteor_score([ans.split()], pred_str.split())
	scores.append(score)
	return sum(scores) / len(scores) if scores else 0.0


	def compute_rouge(preds, answers, answer_vocabs, mapped_classes):
	scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
	scores = []
	for pred, ans, c in zip(preds, answers, mapped_classes):
	if c not in answer_vocabs:
	continue
	pred_idx = pred.argmax(dim=1).item()
	inv_vocab = {v: k for k, v in answer_vocabs[c].items()}
	pred_str = inv_vocab.get(pred_idx, "")
	score = scorer.score(ans, pred_str)["rougeL"].fmeasure
	scores.append(score)
	return sum(scores) / len(scores) if scores else 0.0

	def compute_loss(preds, answers, task_logits, true_q_classes, answer_vocabs,q_types_mapping,q_types,task_heads):
	"""
	preds: list of model predictions for each sample
	answers: list of strings (descriptive answers)
	task_logits: tensor [batch_size, num_task_types]
	true_q_classes: list of lists (fine-grained classes for each question)
	answer_vocabs: dict mapping {q_type: {answer: index}}
	"""

	ce_loss = CrossEntropyLoss()
	mse_loss = MSELoss()

	total_loss = 0

	# 1) Map fine-grained → general classes
	mapped_classes = [
	q_types_mapping[c[0] if isinstance(c, list) else c]
	for c in true_q_classes
	]

	# 2) Question type classification loss
	true_task_types = torch.tensor(
	[q_types.index(c) for c in mapped_classes],
	device=task_logits.device
	)

	#print("task_logits, true_task_types\t",task_logits, true_task_types)

	#print("task_logits, true_task_types\t",task_logits.shape, true_task_types.shape)

	task_loss = ce_loss(task_logits, true_task_types)
	total_loss += task_loss

	# 3) Answer prediction loss (per sample)
	for pred, ans, c in zip(preds, answers, mapped_classes):
	predictor = task_heads[c] # ✅ trained head
	if c == "count":
	# For count, answer must be numeric
	try:
	ans_val = float(ans)
	ans_val = torch.tensor([ans_val], device=pred.device)
	total_loss += mse_loss(pred.squeeze(), ans_val)
	except ValueError:
	print(f"[Warning] Skipping non-numeric count answer: {ans}")
	continue

	else:
	# For categorical tasks (yesno, single, multi, etc.)
	if ans not in answer_vocabs.get(c, {}):
	print(f"[Warning] Skipping unseen or descriptive answer {ans} for task {c}")
	continue

	ans_idx = answer_vocabs[c][ans]

	if ans_idx >= pred.size(1):
	print(f"[Warning] Skipping answer {ans} for task {c}: "
	f"index {ans_idx} >= pred.size(1)")
	continue

	ans_tensor = torch.tensor([ans_idx], device=pred.device)
	total_loss += ce_loss(pred, ans_tensor)

	meteor = compute_meteor(preds, answers, answer_vocabs, mapped_classes)
	print(f"Validation METEOR: {meteor:.4f}")
	#rouge = compute_rouge(preds, answers, answer_vocabs, mapped_classes)
	#print(f"Validation ROUGE-L: {rouge:.4f}")


	return total_loss / len(preds)