Update app.py

app.py

@@ -3,12 +3,18 @@ from transformers import BertTokenizer, BertModel, GPT2LMHeadModel, GPT2Tokenize
 import numpy as np
 import pandas as pd
 import os
-import json
-from fastapi import FastAPI
-from pydantic import BaseModel
+import gradio as gr
 
-app = FastAPI()
+# Load the models and tokenizers
+bert_model_name = 'bert-base-uncased'
+bert_tokenizer = BertTokenizer.from_pretrained(bert_model_name)
+bert_model = BertModel.from_pretrained(bert_model_name)
 
+gpt2_model_name = 'gpt2'
+gpt2_tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)
+gpt2_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name)
+
+# Load the data
 data = {
     "questions": [
         "What is Rookus?",
@@ -33,14 +39,6 @@ data = {
     "default_answers": "I'm sorry, I cannot answer this right now. Your question has been saved, and we will get back to you with a response soon."
 }
 
-bert_model_name = 'models/bert'
-bert_tokenizer = BertTokenizer.from_pretrained(bert_model_name)
-bert_model = BertModel.from_pretrained(bert_model_name)
-
-gpt2_model_name = 'models/gpt2'
-gpt2_tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)
-gpt2_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name)
-
 def get_bert_embeddings(texts):
     inputs = bert_tokenizer(texts, return_tensors='pt', padding=True, truncation=True)
     with torch.no_grad():
@@ -66,12 +64,7 @@ def generate_gpt2_response(prompt, model, tokenizer, max_length=100):
     outputs = model.generate(inputs, max_length=max_length, num_return_sequences=1)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-class QueryRequest(BaseModel):
-    query: str
-
-@app.post("/query/")
-def answer_query(request: QueryRequest):
-    user_query = request.query
+def answer_query(user_query):
     closest_question, similarity = get_closest_question(user_query, data['questions'], threshold=0.95)
     if closest_question and similarity >= 0.95: 
         answer_index = data['questions'].index(closest_question)
@@ -89,8 +82,15 @@ def answer_query(request: QueryRequest):
             df.to_excel(writer, index=False)
         answer = data['default_answers']
 
-    return {"query": user_query, "answer": answer}
+    return answer
+
+iface = gr.Interface(
+    fn=answer_query,
+    inputs="text",
+    outputs="text",
+    title="Rookus AI Query Interface",
+    description="Ask questions about Rookus and get answers generated by AI."
+)
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    iface.launch()