Rename application.py to app.py

application.py → app.py

@@ -6,10 +6,11 @@ from flask import Flask, jsonify, request
 import torch.nn.functional as F
 import boto3
 import pandas as pd
-bucket = 'data-ai-dev2'
+#bucket = 'data-ai-dev2'
 from transformers import BertTokenizer, BertModel
 from torch import cuda
-device = 'cuda' if cuda.is_available() else 'cpu'
+import gradio as gr
+#device = 'cuda' if cuda.is_available() else 'cpu'
 
 class RobertaClass(torch.nn.Module):
     def __init__(self):
@@ -32,10 +33,7 @@ class RobertaClass(torch.nn.Module):
 model = RobertaClass()
 model.to(device)    
 
-s3 = boto3.client('s3', aws_access_key_id='AKIAW5BGUY6ZRCSQBSIJ', aws_secret_access_key= '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')
-s3.download_file(Bucket=bucket, Key='model_hf/tweet_model/tweet_model_v1.bin', Filename = './tweet_model_v1.bin')
-
-model = torch.load('tweet_model_v1.bin', map_location=torch.device('cpu'))
+model = torch.load('./tweet_model_v1.bin', map_location=torch.device('cpu'))
 
 tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased', truncation=True, do_lower_case=True)
 
@@ -43,24 +41,11 @@ def id2class_fun(lst, map_cl):
     s = pd.Series(lst)
     return s.map(map_cl).tolist()
     
-application = Flask(__name__)
-@application.route('/')
-def home():
-    return "Working!"
-
-@application.route('/process/', methods=['POST'])
-def process():
-    content_type = request.headers.get('Content-Type')
-    if (content_type == 'application/json'):
-        json_file = request.json
-    loaded = json.dumps(json_file)
-    new_loaded = json.loads(loaded)
-    text = new_loaded['text']
-    id2class = {0: 'InappropriateUndesirable', 1 : 'GreenContent', 2 : 'IllegalActivities', 
+id2class = {0: 'InappropriateUndesirable', 1 : 'GreenContent', 2 : 'IllegalActivities', 
                 3 : 'DiscriminatoryHate', 4 :'ViolentGraphic', 5:'PotentialAddiction', 
                 6 : 'ExtremismTerrorism', 7 : 'SexualExplicit'}
+def process(text):
     try:
-
         inputs = (
             tokenizer.encode_plus(
                 text, None, add_special_tokens=True, max_length = 512,
@@ -78,14 +63,13 @@ def process():
         prd_score = top_values.cpu().detach().numpy().tolist()
         prd_score = [item for sublist in prd_score for item in sublist]
         otp = dict(zip(prd_cls_1, prd_score))
-    #     .replace(map_class, inplace=True)
-        return jsonify({'output':otp})
+        return {'output':otp}
     except:
-        return jsonify({'output':'something went wrong'})
+        return {'output':'something went wrong'}
+
+inputs = [gr.inputs.Textbox(lines=2, label="Enter the tweet")]
+outputs = gr.outputs.Textbox(label="result")
 
-     
-        
-if __name__ == "__main__":
-    application.debug = True
-    application.run()
+gr.Interface(fn=process, inputs=inputs, outputs=outputs, title="twitter_classifier",
+             theme="compact").launch()