Update app.py

app.py

@@ -1,70 +1,73 @@
-# Import required libraries
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+import torch
 import json
 
-# Create FastAPI app instance
-app = FastAPI()
-
-class ChatBot:
-    def __init__(self):
-        # Load DialoGPT model and tokenizer
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
-            self.model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Model loading failed: {e}")
-
-        # Load courses data
-        try:
-            with open("uts_courses.json", "r") as file:
-                self.courses_data = json.load(file)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Courses data loading failed: {e}")
-
-    def generate_response(self, user_input: str):
-        """
-        Generate response based on user input
-        Args:
-            user_input: User input text
-        Returns:
-            Generated response text
-        """
-        if user_input.lower() == "help":
-            return "I can help you with UTS courses information, feel free to ask!"
-        elif user_input.lower() == "exit":
-            return "Goodbye!"
-        elif user_input.lower() == "list courses":
-            # Generate course list
-            course_list = "\n".join([f"{category}: {', '.join(courses)}" for category, courses in self.courses_data["courses"].items()])
-            return f"Here are the available courses:\n{course_list}"
-        elif user_input.lower() in self.courses_data["courses"]:
-            # List courses under the specified course category
-            return f"The courses in {user_input} category are: {', '.join(self.courses_data['courses'][user_input])}"
-        else:
-            # Use DialoGPT model to generate response
-            input_ids = self.tokenizer.encode(user_input + self.tokenizer.eos_token, return_tensors="pt")
-            response_ids = self.model.generate(input_ids, max_length=100, pad_token_id=self.tokenizer.eos_token_id)
-            response = self.tokenizer.decode(response_ids[0], skip_special_tokens=True)
-            return response
-
-# Define user input model
-class UserInput(BaseModel):
-    user_input: str
-
-# Create chatbot instance
-chatbot = ChatBot()
-
-# Define API route
-@app.post("/")
-async def chat(user_input: UserInput):
-    """
-    Process user input and return response
-    Args:
-        user_input: User input JSON data
-    Returns:
-        JSON data containing the response text
-    """
-    response = chatbot.generate_response(user_input.user_input)
-    return {"response": response}
+
+title = "????AI ChatBot"
+description = "A State-of-the-Art Large-scale Pretrained Response generation model (DialoGPT)"
+examples = [["What courses are available in Engineering?"]]
+
+
+tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
+model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model.to(device)
+
+# Load courses data from JSON file
+with open("uts_courses.json", "r") as f:
+    courses_data = json.load(f)
+
+
+def predict(input, history=[]):
+    # Check if the input question is about courses
+    if "courses" in input.lower():
+        # Check if the input question contains a specific field (e.g., Engineering, Information Technology, etc.)
+        for field in courses_data["courses"]:
+            if field.lower() in input.lower():
+                # Get the list of courses for the specified field
+                courses_list = courses_data["courses"][field]
+                # Format the response
+                response = f"The available courses in {field} are: {', '.join(courses_list)}."
+                return response, history
+
+    # If the input question is not about courses, use the dialogue model to generate a response
+    # tokenize the new input sentence
+    new_user_input_ids = tokenizer.encode(
+        input + tokenizer.eos_token, return_tensors="pt"
+    ).to(device)
+
+    # append the new user input tokens to the chat history
+    bot_input_ids = torch.cat([torch.tensor(history).to(device), new_user_input_ids], dim=-1)
+
+    # generate a response
+    history = model.generate(
+        bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id
+    ).tolist()
+
+    # convert the tokens to text, and then split the responses into lines
+    response = tokenizer.decode(history[0]).split("")
+    response = [
+        (response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)
+    ]  # convert to tuples of list
+    return response, history
+
+
+def main():
+    pass
+
+
+if __name__ == "__main__":
+    main()
+
+
+gr.Interface(
+    fn=predict,
+    title=title,
+    description=description,
+    examples=examples,
+    inputs=["text", "state"],
+    outputs=["chatbot", "state"],
+    theme="finlaymacklon/boxy_violet",
+).launch()