Create app.py

app.py

@@ -0,0 +1,79 @@
+import streamlit as st
+from datasets import load_dataset
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+import torch
+
+ds = load_dataset("higgsfield/school-math-questions")
+
+class MathDataset(torch.utils.data.Dataset):
+    def __init__(self, qa_pairs, tokenizer, max_length=128):
+        self.qa_pairs = qa_pairs
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self):
+        return len(self.qa_pairs)
+
+    def __getitem__(self, idx):
+        question, answer = self.qa_pairs[idx]
+        input_text = f"Q: {question} A:"
+        
+        # Tokenize and pad input and target sequences
+        input_ids = self.tokenizer.encode(input_text, truncation=True, padding="max_length", max_length=self.max_length, return_tensors="pt").squeeze(0)
+        target_ids = self.tokenizer.encode(answer.strip(), truncation=True, padding="max_length", max_length=self.max_length, return_tensors="pt").squeeze(0)
+
+        # Set the labels to -100 where input_ids are padding tokens
+        target_ids[target_ids == self.tokenizer.pad_token_id] = -100
+        
+        return {
+            "input_ids": input_ids,
+            "labels": target_ids,
+        }
+
+model_name = "gpt2"
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
+model = GPT2LMHeadModel.from_pretrained(model_name)
+
+math_dataset = MathDataset(qa_pairs, tokenizer)
+
+from transformers import Trainer, TrainingArguments
+
+# Set training arguments
+training_args = TrainingArguments(
+    output_dir="./results",
+    num_train_epochs=3,
+    per_device_train_batch_size=2,
+    save_steps=10,
+    save_total_limit=2,
+)
+
+# Create a Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=math_dataset,
+)
+
+# Fine-tune the model
+trainer.train()
+
+class MathChatBot:
+    def __init__(self, model_name="gpt2"):
+        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+        self.model = GPT2LMHeadModel.from_pretrained(model_name)
+
+    def get_response(self, question):
+        input_text = f"Q: {question} A:"
+        input_ids = self.tokenizer.encode(input_text, return_tensors="pt")
+
+        output = self.model.generate(input_ids, max_length=50, num_return_sequences=1)
+        answer = self.tokenizer.decode(output[0], skip_special_tokens=True)
+        return answer.split("A:")[-1].strip()
+
+# Usage
+if __name__ == "__main__":
+    bot = MathChatBot()
+    user_input = st.text_area("Enter your question:"
+    response = bot.get_response(user_input)
+    st.write(f"Bot: {response}")