Update app.py

app.py

@@ -1,12 +1,28 @@
+from flask import Flask, request, jsonify
+from transformers import DistilBertTokenizerFast, TFDistilBertForSequenceClassification
+import tensorflow as tf
+import numpy as np
+import torch
 
-import gradio as gr
-from transformers import pipeline
+app = Flask(__name__)
+device = '/GPU:0' if torch.cuda.is_available() else 'CPU'
+model_name = "distilbert-base-uncased"
+tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
+model = TFDistilBertForSequenceClassification.from_pretrained(model_name, from_pt=True).signatures['serving_default'].to(device)
+session = tf.compat.v1.keras.backend.get_session()
 
-dialogpt = pipeline("text-generation", model="microsoft/DialoGPT-medium")
+@app.route("/predict", methods=["POST"])
+def predict():
+    data = request.get_json()
+    input_text = data["input"]
+    input_ids = torch.tensor(tokenizer.encode(input_text)).unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        outputs = model(inputs={'input_ids': input_ids}).logits
+        probabilities = tf.nn.softmax(outputs).numpy()
+        prediction = np.argmax(probabilities)
+        
+    return jsonify({"response": prediction})
 
-def respond(context):
-    response = dialogpt(context, max_length=100, do_sample=True)
-    return response[0]['generated_text'][len(context):]
-
-iface = gr.Interface(fn=respond, inputs="text", outputs="text")
-iface.launch()
+if __name__ == "__main__":
+    app.run(debug=True)