Update app.py

app.py

@@ -1,29 +1,55 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from huggingface_hub import Repository
+import os
 
-# Load the fine-tuned model
-model_name = "./tuned_model"  # Load from Hugging Face or locally
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Define model directory and Hugging Face repo name
+model_dir = "./tuned_model"  # Directory where the fine-tuned model is saved
+repo_name = "krisha06/Python_tutor"  # Replace with your Hugging Face repo name
 
-# Set up Streamlit UI
+# Streamlit App Title
 st.title("AI Coding Mentor")
-st.write("Ask me any programming-related question!")
 
-# User input (question)
+# Section to upload the fine-tuned model to Hugging Face Hub
+st.header("Upload Your Fine-Tuned Model to Hugging Face Hub")
+
+# Option to upload the model
+upload_model_button = st.button("Upload Model to Hugging Face Hub")
+
+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()
+        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.")
+
+# Section for using the model in the app
+st.header("Ask Me Any Coding Question!")
+
+# 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
+
+# Check if the model is uploaded to Hugging Face Hub
+if os.path.exists(model_dir):
+    # Load the model and tokenizer from local directory
+    model = AutoModelForCausalLM.from_pretrained(model_dir, load_in_8bit=True)
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+else:
+    # Load the model from Hugging Face Hub if it's not found locally
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# User input: Coding question
 question = st.text_input("Enter your coding question:")
 
 if question:
-    # Prepare the input for the model
     input_text = f"### Question:\n{question}\n### Answer:"
     inputs = tokenizer(input_text, return_tensors="pt")
 
-    # Generate the answer using the fine-tuned model
-    output = model.generate(**inputs, max_length=200, num_return_sequences=1)
-
-    # Decode the output
-    answer = tokenizer.decode(output[0], skip_special_tokens=True)
+    with st.spinner("Processing..."):
+        output = model.generate(**inputs, max_length=200, num_return_sequences=1)
+        answer = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    # Display the result
     st.write(answer)