Update app.py

app.py

@@ -1,57 +1,97 @@
+import torch
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import Repository
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForSeq2Seq
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model
 import os
 
-# Retrieve Hugging Face token from Streamlit secrets
-token = st.secrets["HUGGINGFACE_TOKEN"]  # Make sure to add your Hugging Face token to Streamlit secrets
+# UI
+st.title("AI Tutor (Fine-tuned LLM)")
+st.write("This AI tutor is fine-tuned on Python-related questions.")
 
-# Define model directory and Hugging Face repo name
-model_dir = "./tuned_model"  # The directory where the fine-tuned model is saved
-repo_name = "krisha06/Python_tutor"  # Replace with your Hugging Face repo name
+# Load base model and tokenizer
+model_name = "microsoft/phi-2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Streamlit App Title
-st.title("AI Coding Mentor")
+# 🔥 Fix: Add padding token
+tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
-# Section to upload the fine-tuned model to Hugging Face Hub
-st.header("Upload Your Fine-Tuned Model to Hugging Face Hub")
+# Check if fine-tuned model exists
+model_path = "./models"
+if os.path.exists(model_path):
+    st.write("✅ Loading fine-tuned model...")
+    model = AutoModelForCausalLM.from_pretrained(model_path)
+else:
+    st.write("⚡ Fine-tuning the model (this will take time)...")
 
-# Option to upload the model
-upload_model_button = st.button("Upload Model to Hugging Face Hub")
+    # Load model on CPU
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.bfloat16,
+        device_map={"": "cpu"}  # Force CPU usage
+    )
 
-if upload_model_button:
-    if os.path.exists(model_dir):
-        # Initialize the Hugging Face Repository and push model to the Hub
-        repo = Repository(local_dir=model_dir, clone_from=repo_name)
-        repo.push_to_hub(token=token)  # Use the token for authentication
-        st.success("Model uploaded to Hugging Face Hub successfully!")
-    else:
-        st.error("Model directory does not exist. Please make sure the model is fine-tuned first.")
+    # Resize model embeddings
+    model.resize_token_embeddings(len(tokenizer))
 
-# Section for using the model in the app
-st.header("Ask Me Any Coding Question!!")
+    # Apply LoRA
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=32,
+        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    model = get_peft_model(model, lora_config)
 
-# Load model and tokenizer (either from local directory or Hugging Face Hub)
-model_name = repo_name  # Use the repo name if the model is on Hugging Face Hub, else use local dir
+    # Load dataset and tokenize
+    dataset = load_dataset("mbpp", split="train")
 
-if os.path.exists(model_dir):
-    # Load the model and tokenizer from local directory
-    model = AutoModelForCausalLM.from_pretrained(model_dir, use_auth_token=token)
-    tokenizer = AutoTokenizer.from_pretrained(model_dir, use_auth_token=token)
-else:
-    # Load the model from Hugging Face Hub
-    model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=token)
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
+
+    tokenized_dataset = dataset.map(tokenize_function, batched=True)
+
+    # Data collator
+    data_collator = DataCollatorForSeq2Seq(tokenizer, return_tensors="pt")
+
+    # Training arguments
+    training_args = TrainingArguments(
+        per_device_train_batch_size=1,
+        per_device_eval_batch_size=1,
+        num_train_epochs=1,  # Reduce epochs for quick training
+        learning_rate=3e-4,
+        output_dir=model_path,
+        save_strategy="epoch",
+        logging_dir="./logs",
+        logging_steps=10,
+        evaluation_strategy="epoch",
+        save_total_limit=2,
+        load_best_model_at_end=True
+    )
+
+    # Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+        data_collator=data_collator,
+    )
 
-# User input: Coding question
-question = st.text_input("Enter your coding question:")
+    # Train
+    trainer.train()
 
-if question:
-    input_text = f"### Question:\n{question}\n### Answer:"
-    inputs = tokenizer(input_text, return_tensors="pt")
+    # Save model
+    model.save_pretrained(model_path)
+    tokenizer.save_pretrained("./tokenizer")
 
-    with st.spinner("Processing..."):
-        output = model.generate(**inputs, max_length=200, num_return_sequences=1)
-        answer = tokenizer.decode(output[0], skip_special_tokens=True)
+    st.write("🎉 Fine-tuning complete! Model saved.")
 
-    st.write(answer)
+# Chat Interface
+user_input = st.text_input("Ask a coding question:")
+if user_input:
+    inputs = tokenizer(user_input, return_tensors="pt").to("cpu")
+    outputs = model.generate(**inputs, max_length=150)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    st.write("🤖 AI Tutor:", response)