Create chat.py

chat.py

@@ -0,0 +1,118 @@
+import os
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+import torch
+from torch.utils.data import Dataset
+
+# Initialize model and tokenizer as global variables
+model = None
+tokenizer = None
+
+# Dictionary to store user instructions for future responses
+user_instructions = {}
+
+# Dummy dataset class for user feedback
+class FeedbackDataset(Dataset):
+    def __init__(self, input_texts, target_texts):
+        self.input_texts = input_texts
+        self.target_texts = target_texts
+
+    def __len__(self):
+        return len(self.input_texts)
+
+    def __getitem__(self, idx):
+        inputs = tokenizer.encode(self.input_texts[idx], return_tensors="pt").squeeze()
+        targets = tokenizer.encode(self.target_texts[idx], return_tensors="pt").squeeze()
+        return {"input_ids": inputs, "labels": targets}
+
+def load_model(model_name_or_path):
+    global model, tokenizer
+
+    st.write(f"Loading model from {model_name_or_path}...")
+
+    # Load the tokenizer and model
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
+
+    st.success("Model loaded successfully!")
+
+def generate_response(input_text):
+    # Ensure model and tokenizer are loaded
+    if model is None or tokenizer is None:
+        st.error("Model is not loaded. Please load a model first.")
+        return ""
+
+    # Check if there's a user-defined response
+    if input_text in user_instructions:
+        return user_instructions[input_text]
+
+    # Encode input text
+    inputs = tokenizer.encode(input_text, return_tensors="pt")
+
+    # Generate response using the model
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs, max_length=100, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95
+        )
+
+    # Decode and return the response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+def train_on_feedback(input_text, correct_response):
+    # Prepare dataset
+    dataset = FeedbackDataset([input_text], [correct_response])
+
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir="./feedback_model",
+        num_train_epochs=1,
+        per_device_train_batch_size=1,
+        learning_rate=1e-5,
+        logging_dir='./logs',
+        logging_steps=10,
+        save_steps=100
+    )
+
+    # Trainer for the feedback loop
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset,
+    )
+
+    # Train model on the feedback
+    trainer.train()
+
+def chat_interface():
+    st.title("🤖 Chat with AI")
+
+    # Input for model name or path
+    model_name_or_path = st.text_input("Enter model name or local path:", "gpt2")
+
+    # Button to load the model
+    if st.button("Load Model"):
+        load_model(model_name_or_path)
+
+    st.write("---")
+
+    # Chat input
+    input_text = st.text_input("You:")
+
+    if st.button("Send"):
+        response = generate_response(input_text)
+        st.write("AI:", response)
+
+        # Feedback section
+        feedback = st.radio("Was this response helpful?", ("Yes", "No"))
+
+        if feedback == "No":
+            correct_response = st.text_input("What should the AI have said?")
+            if st.button("Submit Feedback"):
+                # Train model on feedback
+                train_on_feedback(input_text, correct_response)
+                st.success("Feedback recorded. AI will improve based on this feedback.")
+
+# Run chat interface
+if __name__ == "__main__":
+    chat_interface()