Update app.py

app.py

@@ -1,79 +1,51 @@
 import gradio as gr
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModel
-import dask.dataframe as dd
-from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load models and tokenizer
+# Load models and tokenizers
 def load_models():
-    # Load model 1
-    model_1 = AutoModel.from_pretrained("Canstralian/RedTeamAI")
-    
-    # Load model 2
-    model_2 = AutoModel.from_pretrained("mradermacher/BashCopilot-6B-preview-GGUF")
-    
-    # Load tokenizer and sequence classification model
-    tokenizer = AutoTokenizer.from_pretrained("bash1130/bert-base-finetuned-ynat")
-    model_3 = AutoModelForSequenceClassification.from_pretrained("bash1130/bert-base-finetuned-ynat")
-    
-    return model_1, model_2, tokenizer, model_3
+    # Load a conversational model and tokenizer (you can customize it further)
+    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+    return model, tokenizer
 
-# Load dataset using Dask
-def load_data():
-    # Example of loading a dataset using Dask (adjust paths as necessary)
-    splits = {'creative_content': 'data/creative_content-00000-of-00001.parquet'}
-    df = dd.read_parquet("hf://datasets/microsoft/orca-agentinstruct-1M-v1/" + splits["creative_content"])
-    return df.head()
+# Generate responses
+def chat_with_model(user_input, model, tokenizer, chat_history):
+    # Tokenize the user input and chat history
+    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
 
-# Function for model inference
-def infer_model(input_text, model_type):
-    # Choose the model based on the input (you can add more models or conditions as needed)
-    if model_type == 'RedTeamAI':
-        model = models[0]
-    elif model_type == 'BashCopilot':
-        model = models[1]
-    elif model_type == 'BertModel':
-        model = models[3]
-        inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
-        outputs = model(**inputs)
-        return outputs.logits.argmax(dim=-1).item()
-    else:
-        return "Model type not recognized."
-    
-    # If you need to generate outputs based on the models directly, you can use:
-    # outputs = model.generate(input_text) or other inference methods depending on the model.
-    return f"Model {model_type} inference not implemented yet."
+    # Append new user input to chat history
+    bot_input_ids = torch.cat([chat_history, new_user_input_ids], dim=-1) if chat_history is not None else new_user_input_ids
 
-# Gradio Interface setup
-def build_interface():
-    # Load models and data
-    model_1, model_2, tokenizer, model_3 = load_models()
-    global models
-    models = [model_1, model_2, tokenizer, model_3]
-    
-    # Load the dataset (example function, you can add more functionality)
-    data_preview = load_data()
-    
-    print(f"Dataset preview: {data_preview}")
-    
-    # Create Gradio interface
+    # Generate a response from the model
+    chat_history = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+
+    # Decode the model's output and return
+    bot_output = tokenizer.decode(chat_history[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+    return chat_history, bot_output
+
+# Initialize model and tokenizer
+model, tokenizer = load_models()
+
+# Build Gradio interface
+def build_gradio_interface():
     with gr.Blocks() as demo:
-        gr.Markdown("# Chagrin AI - Model Inference & Dataset Explorer")
-        
-        # Model selection dropdown
-        model_type = gr.Dropdown(choices=["RedTeamAI", "BashCopilot", "BertModel"], label="Choose Model")
-        
-        # Textbox for user input
-        input_text = gr.Textbox(label="Enter your input text")
-        
-        # Button to trigger inference
-        result = gr.Textbox(label="Inference Result")
-        
-        submit_btn = gr.Button("Run Inference")
-        submit_btn.click(infer_model, inputs=[input_text, model_type], outputs=result)
-    
+        gr.Markdown("# Chagrin AI Chatbot")
+
+        # Set up chat window
+        chatbot = gr.Chatbot()
+
+        # Create text input box for user to type
+        user_input = gr.Textbox(label="Type your message", placeholder="Ask something...", interactive=True)
+
+        # Create button for sending the input
+        submit_btn = gr.Button("Send Message")
+
+        # Button click function
+        submit_btn.click(chat_with_model, inputs=[user_input, model, tokenizer, chatbot], outputs=[chatbot, chatbot])
+
     demo.launch()
 
-# Run the app
+# Run the Gradio interface
 if __name__ == "__main__":
-    build_interface()
+    build_gradio_interface()