Update app.py

app.py

@@ -1,69 +1,30 @@
-import torch
+import os
+import subprocess
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
-from datasets import load_dataset
 
-st.set_page_config(page_title="AstroChat", layout="centered")
+# ✅ Manually install PyTorch if not found
+try:
+    import torch
+except ModuleNotFoundError:
+    st.warning("Installing missing dependencies... Please wait.")
+    subprocess.run(["pip", "install", "torch", "torchvision", "torchaudio"], check=True)
+    import torch  # Re-import after installation
 
-# ✅ Load dataset (WikiText)
-dataset = load_dataset("wikitext", "wikitext-103-raw-v1")
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# ✅ Load tokenizer & model (DistilGPT-2)
+# ✅ Load tokenizer & model
 model_name = "distilgpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# ✅ Ensure tokenizer has a padding token
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
-# ✅ Tokenization function
-def tokenize_function(examples):
-    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
-
-# ✅ Tokenize dataset
-tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
-
-# ✅ Detect GPU automatically
+# ✅ Detect GPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
-
-# ✅ Define training arguments
-training_args = TrainingArguments(
-    output_dir="./results",
-    evaluation_strategy="epoch",
-    per_device_train_batch_size=8,
-    per_device_eval_batch_size=8,
-    num_train_epochs=3,
-    save_strategy="epoch",
-    save_total_limit=2,
-    load_best_model_at_end=True,
-    logging_dir="./logs",
-    logging_steps=500,
-    fp16=True if device == "cuda" else False,  # Use mixed precision on GPU
-    no_cuda=not torch.cuda.is_available(),
-    push_to_hub=True  # ✅ Upload model to Hugging Face Hub
-)
-
-# ✅ Trainer setup
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_datasets["train"],
-    eval_dataset=tokenized_datasets["validation"],
-)
-
-# ✅ Train the model
-if __name__ == "__main__":
-    trainer.train()
-    trainer.save_model("./trained_model")  # Save trained model
-    tokenizer.save_pretrained("./trained_model")
+model.to(device)
 
 # ✅ Text Generation Function
 def generate_text(prompt, max_length=150, temperature=0.8, top_p=0.9):
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
-    model.to(device)
-
     output = model.generate(
         input_ids,
         max_length=max_length,
@@ -75,7 +36,7 @@ def generate_text(prompt, max_length=150, temperature=0.8, top_p=0.9):
 
 # ✅ Streamlit UI
 st.title("🔮 AstroChat - AI Chatbot")
-st.write("A conversational AI trained on WikiText using DistilGPT-2.")
+st.write("A conversational AI trained using DistilGPT-2.")
 
 user_input = st.text_input("Ask me anything:")
 if st.button("Generate Response"):