Update QA_Tensorflow.py

QA_Tensorflow.py

@@ -40,20 +40,24 @@ def question_answering_tf(question, context):
 
     #Finding (start token, end token) pair with best probability score
     max_score = 0.0
-    start_index,end_index = 0,0
+    start_token = 0
+    end_token = 0
+    offset_index = 0
+    
     for i, probs in enumerate(zip(start_probabilities, end_probabilities)):
       sp, ep = probs
-      scores = tf.matmul(sp[:,np.newaxis], ep[np.newaxis,:])
+      scores = sp[:,np.newaxis] * ep[np.newaxis,:]
       index = np.triu(scores).argmax().item()
-      row = int(index/384)
-      col = index % 384
+      row = index // scores.shape[1]
+      col = index % scores.shape[1]
       score = scores[row][col]
       if(score > max_score):
         max_score = score
-        start_index = row
-        end_index = col
+        start_token = row
+        end_token = col
+        offset_index = i
 
     #Return characters from context corresponding to start and end of token characters 
-    start = int(offset_mapping[i][start_index][0])
-    end = int(offset_mapping[i][end_index][1])
-    return context[start:end+1]
+    start_char, _ = offset_mapping[offset_index][start_token]
+    _, end_char = offset_mapping[offset_index][end_token]
+    return context[start_char:end_char+1]