Update app.py

app.py

@@ -1,28 +1,40 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import sentencepiece
+import gradio as gr
+import re
 import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("ahmed792002/Finetuning_T5_HealthCare_Chatbot", use_fast=True)
-model = AutoModelForSeq2SeqLM.from_pretrained("ahmed792002/Finetuning_T5_HealthCare_Chatbot")
+tokenizer = T5Tokenizer.from_pretrained("ahmed792002/Finetuning_T5_HealthCare_Chatbot")
+model = T5ForConditionalGeneration.from_pretrained("ahmed792002/Finetuning_T5_HealthCare_Chatbot")
 
-# Define the chatbot function
-def chatbot(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
-    outputs = model.generate(inputs["input_ids"], max_length=100, num_beams=4, early_stopping=True)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
-
-# Set up the Gradio interface
-interface = gr.Interface(
-    fn=chatbot,
-    inputs=gr.Textbox(label="Enter your query"),  # Corrected import
-    outputs=gr.Textbox(label="Response"),         # Corrected import
-    title="Healthcare Chatbot",
-    description="Ask healthcare-related questions, and get responses from the fine-tuned T5 model."
+def clean_text(text):
+    text = re.sub(r'\r\n', ' ', text)  # Remove carriage returns and line breaks
+    text = re.sub(r'\s+', ' ', text)  # Remove extra spaces
+    text = re.sub(r'<.*?>', '', text)  # Remove any XML tags
+    text = text.strip().lower()  # Strip and convert to lower case
+    return text
+def chatbot(query):
+    query = clean_text(query)
+    input_ids = tokenizer(query,return_tensors="pt",max_length=256,truncation=True)
+    inputs = {key: value.to(device) for key, value in input_ids.items()}
+    outputs = model.generate(
+        input_ids["input_ids"],
+        max_length=1024,
+        num_beams=5,
+        early_stopping=True
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+	
+"""
+For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+"""
+demo = gr.ChatInterface(
+    chatbot,
+    additional_inputs=[
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+    ],
 )
 
-# Launch the app
+
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()