Google Colab imports

scripts.py

@@ -20,7 +20,7 @@ import seaborn
 import tqdm
 import gradio
 import boto3
-import google.colab
+from google.colab import userdata
 
 class SequenceCrossEntropyLoss(torch.nn.Module):
     def __init__(self):
@@ -61,8 +61,8 @@ def download_training_data():
     if not os.path.exists('corpora'):
         os.makedirs('corpora')
     s3 = boto3.client('s3',
-                      aws_access_key_id=google.colab.userdata.get('AWS_KEY'),
-                      aws_secret_access_key=google.colab.userdata.get('AWS_SECRET'))
+                      aws_access_key_id=userdata.get('AWS_KEY'),
+                      aws_secret_access_key=userdata.get('AWS_SECRET'))
     for obj in s3.list_objects(Bucket='qarac')['Contents']:
         filename = obj['Key']
         s3.download_file('qarac',filename,'corpora/{}'.format(filename))
@@ -155,7 +155,7 @@ def train_models(path,progress=gradio.Progress(track_tqdm=True)):
                 epoch_data[batch] = loss.item()
         history[epoch_label] = epoch_data
         scheduler.step()
-    huggingface_hub.login(token=google.colab.userdata.get('HUGGINGFACE_TOKEN'))
+    huggingface_hub.login(token=userdata.get('HUGGINGFACE_TOKEN'))
     trainer.question_encoder.push_to_hub('{}/qarac-roberta-question-encoder'.format(path))
     trainer.answer_encoder.push_to_hub('{}/qarac-roberta-answer-encoder'.format(path))
     trainer.decoder.push_to_hub('{}/qarac-roberta-decoder'.format(path))