Update app.py

app.py

@@ -1,45 +1,55 @@
-import sentencepiece
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-import re
 import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration
 
 # Load pre-trained model and tokenizer
-tokenizer = T5Tokenizer.from_pretrained("ahmed792002/Finetuning_T5_HealthCare_Chatbot")
-model = T5ForConditionalGeneration.from_pretrained("ahmed792002/Finetuning_T5_HealthCare_Chatbot")
+tokenizer = AutoTokenizer.from_pretrained("ahmed792002/alzheimers_memory_support_ai")
+model = AutoModelForCausalLM.from_pretrained("ahmed792002/alzheimers_memory_support_ai")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)  # Send the model to the correct device
 
 # Function to clean input text
 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
+    return text.strip()  # Simply remove leading/trailing spaces
 
 # Chatbot function
-def chatbot(query, history, system_message):
+def chatbot(query, history, system_message, max_length, temperature, top_k, top_p):
+    """
+    Processes a user query through the specified model to generate a response.
+    """
+    # Clean the input 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,  # Adjust this as needed for your use case
-        num_beams=5,
-        temperature=0.7,  # Adjust this as needed
-        top_p=0.95,  # Adjust this as needed
-        early_stopping=True
+    
+    # Tokenize input query
+    input_ids = tokenizer.encode(query, return_tensors="pt").to(device)
+    
+    # Generate text using the model
+    final_outputs = model.generate(
+        input_ids,
+        do_sample=True,
+        max_length=int(max_length),  # Convert max_length to integer
+        temperature=float(temperature),  # Convert temperature to float
+        top_k=int(top_k),  # Convert top_k to integer
+        top_p=float(top_p),  # Convert top_p to float
+        pad_token_id=tokenizer.pad_token_id,
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Decode generated text
+    response = tokenizer.decode(final_outputs[0], skip_special_tokens=True)
+    return response
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+# Gradio ChatInterface
 demo = gr.ChatInterface(
     chatbot,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(value="You are a friendly chatbot.", label="System message"),
+        gr.Slider(128, 1024, value=256, step=64, label="Max Length"),  # Slider for max_length
+        gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature"),  # Slider for temperature
+        gr.Slider(1, 100, value=50, step=1, label="Top-K"),  # Slider for top_k
+        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-P"),  # Slider for top_p
     ],
+    title="Custom Alzheimer's Memory Support AI",
+    description="This chatbot uses the fine-tuned model 'ahmed792002/alzheimers_memory_support_ai'. Customize settings like max length, temperature, top-k, and top-p for better results.",
 )
 
 if __name__ == "__main__":