Spaces:
Sleeping
Sleeping
File size: 4,939 Bytes
bf8c161 aef03a7 cc4d3bc 96836c8 bf8c161 aef03a7 bf8c161 |
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 |
import torch
import torch.nn as nn
from PIL import Image
import numpy as np
def simple_process_input(image, text_input, model, src_lang, tgt_lang, cfg):
"""Simplified inference that bypasses complex text processing"""
device = next(model.parameters()).device
# Transform image
import datasets.diagram_aug as T_diagram
diagram_transform = T_diagram.Compose([
T_diagram.Resize(cfg.diagram_size),
T_diagram.CenterCrop(cfg.diagram_size),
T_diagram.ToTensor(),
T_diagram.Normalize()
])
diagram = diagram_transform(image).unsqueeze(0).to(device)
# Simple text processing - just tokenize the words
words = text_input.split() if text_input.strip() else ["problem"]
# Map words to indices (use PAD for unknown)
text_indices = []
for word in words:
if word in src_lang.word2index:
text_indices.append(src_lang.word2index[word])
else:
text_indices.append(0) # PAD token
# Ensure minimum length
if len(text_indices) == 0:
text_indices = [0]
# Create text tensors with proper shape
batch_size = 1
text_len = len(text_indices)
# For MLM pretrain, tokens need to be 3D: [batch, seq_len, vocab_size]
# But here we use 2D: [batch, seq_len] and let the embedding layer handle it
token_tensor = torch.LongTensor([text_indices]).to(device)
# Ensure sect_tag and class_tag match token length
sect_tag_indices = [1] * text_len # Default to [PROB]
class_tag_indices = [1] * text_len # Default to [GEN]
# The model expects token to be [batch, num_subwords_per_token, seq_len]
# For simple case, we have 1 subword per token, so shape is [batch, 1, seq_len]
# This gets embedded and summed over dim=1 to get [batch, seq_len, embed_dim]
# Create 3D tensor: [batch_size, 1, text_len]
# Each token is a single subword, so middle dimension is 1
token_tensor_3d = token_tensor.unsqueeze(1) # [batch, 1, seq_len]
text_dict = {
'token': token_tensor_3d,
'sect_tag': torch.LongTensor([sect_tag_indices]).to(device),
'class_tag': torch.LongTensor([class_tag_indices]).to(device),
'len': torch.LongTensor([text_len]).to(device)
}
# Simple var dict (no variables detected)
# Note: var positions need to account for the diagram token that will be added
var_dict = {
'pos': torch.zeros(batch_size, 1, dtype=torch.long).to(device),
'len': torch.zeros(batch_size, dtype=torch.long).to(device),
'var_value': [],
'arg_value': []
}
# Expression dict for inference
exp_dict = {
'exp': torch.LongTensor([[1]]).to(device), # SOS token
'len': torch.ones(batch_size, dtype=torch.long).to(device),
'answer': 0
}
# Run inference with no_grad
with torch.no_grad():
try:
# Create a copy of text_dict to avoid in-place modification
text_dict_copy = {
'token': text_dict['token'].clone(),
'sect_tag': text_dict['sect_tag'].clone(),
'class_tag': text_dict['class_tag'].clone(),
'len': text_dict['len'].clone()
}
var_dict_copy = {
'pos': var_dict['pos'].clone(),
'len': var_dict['len'].clone(),
'var_value': var_dict['var_value'],
'arg_value': var_dict['arg_value']
}
outputs = model(diagram, text_dict_copy, var_dict_copy, exp_dict, is_train=False)
except Exception as e:
return f"Model inference error: {str(e)}"
# Decode outputs
if outputs is not None:
try:
# Handle different output types
if isinstance(outputs, tuple):
outputs = outputs[0]
if isinstance(outputs, torch.Tensor):
if outputs.dim() > 1:
output_indices = outputs[0].cpu().numpy()
else:
output_indices = outputs.cpu().numpy()
else:
output_indices = outputs
# Convert indices to symbols
output_symbols = []
for idx in output_indices:
if idx < len(tgt_lang.index2word):
symbol = tgt_lang.index2word[idx]
if symbol in ['[EOS]', '[PAD]']:
break
if symbol not in ['[SOS]']:
output_symbols.append(symbol)
if output_symbols:
return f"Generated expression: {' '.join(output_symbols)}"
else:
return "No solution generated (empty output)"
except Exception as e:
return f"Output decoding error: {str(e)}"
return "No output from model" |