Update app.py

app.py

@@ -1,6 +1,5 @@
 import json
 import requests
-import os
 import streamlit as st
 import pandas as pd
 import openpyxl
@@ -28,17 +27,24 @@ uploaded_files = {
     "Individual Timetable": st.sidebar.file_uploader("Upload Individual Timetable", type=["xlsx", "docx"]),
 }
 
-# ✅ Load Timetable Data (Excel & DOCX)
-def load_timetable(file):
-    if not file:
-        return None
-    if file.name.endswith(".xlsx"):
-        wb = openpyxl.load_workbook(file)
-        sheet = wb.active
-        return [row for row in sheet.iter_rows(values_only=True)]
-    elif file.name.endswith(".docx"):
-        doc = Document(file)
-        return [para.text for para in doc.paragraphs]
+# ✅ Load XLSX File
+def load_xlsx(file):
+    return pd.read_excel(file)
+
+# ✅ Load DOCX File
+def load_docx(file):
+    doc = Document(file)
+    return "\n".join([para.text for para in doc.paragraphs])
+
+# ✅ Save DOCX File
+def save_docx(text):
+    doc = Document()
+    for line in text.split("\n"):
+        doc.add_paragraph(line)
+    output = BytesIO()
+    doc.save(output)
+    output.seek(0)
+    return output
 
 # ✅ Ask TinyLlama AI via API
 def ask_tinyllama_api(query):
@@ -46,8 +52,15 @@ def ask_tinyllama_api(query):
         return "Error: Please enter your API key."
 
     url = "https://api-inference.huggingface.co/v1/chat/completions"
-    headers = {"Authorization": f"Bearer {hf_api_key}", "Content-Type": "application/json"}
-    payload = {"model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "messages": [{"role": "user", "content": query}], "max_tokens": 500}
+    headers = {
+        "Authorization": f"Bearer {hf_api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        "messages": [{"role": "user", "content": query}],
+        "max_tokens": 500
+    }
 
     response = requests.post(url, headers=headers, json=payload)
     if response.status_code == 200:
@@ -55,44 +68,28 @@ def ask_tinyllama_api(query):
     else:
         return f"API Error: {response.status_code} - {response.text}"
 
-# ✅ Auto-Schedule Missing Slots (Excel & DOCX)
+# ✅ Auto-Schedule Missing Slots in Excel
 def auto_schedule(file):
     if not file:
         return "No timetable uploaded."
 
-    if file.name.endswith(".xlsx"):
-        wb = openpyxl.load_workbook(file)
-        sheet = wb.active
-
-        empty_slots = []
-        for row_idx, row in enumerate(sheet.iter_rows(min_row=2, values_only=True), start=2):
-            if None in row or "" in row:
-                empty_slots.append(row_idx)
-
-        for row_idx in empty_slots:
-            query = f"Suggest a subject and faculty for the empty slot in row {row_idx}."
-            suggestion = ask_tinyllama_api(query)
-
-            try:
-                subject, faculty = suggestion.split(", Faculty: ")
-                sheet.cell(row=row_idx, column=4, value=subject.strip())
-                sheet.cell(row=row_idx, column=5, value=faculty.strip())
-            except:
-                continue  
-
-        return f"Auto-scheduling completed for {len(empty_slots)} slots."
-
-    elif file.name.endswith(".docx"):
-        doc = Document(file)
-        new_doc = Document()
-        for para in doc.paragraphs:
-            if "EMPTY SLOT" in para.text:
-                query = "Suggest a subject and faculty for this slot."
-                suggestion = ask_tinyllama_api(query)
-                new_doc.add_paragraph(suggestion)
-            else:
-                new_doc.add_paragraph(para.text)
-        return new_doc
+    df = pd.read_excel(file)
+    empty_slots = df[df.isnull().any(axis=1)]
+
+    for index, row in empty_slots.iterrows():
+        query = f"Suggest a subject and faculty for an empty slot on {row.get('Day', 'Unknown')} at {row.get('Time', 'Unknown')}."
+        suggestion = ask_tinyllama_api(query)
+        try:
+            subject, faculty = suggestion.split(", Faculty: ")
+            df.loc[index, "Subject"] = subject.strip()
+            df.loc[index, "Faculty"] = faculty.strip()
+        except:
+            continue  
+
+    output = BytesIO()
+    df.to_excel(output, index=False, engine='openpyxl')
+    output.seek(0)
+    return output
 
 # ✅ AI Query Section
 st.markdown("## 🤖 Ask TinyLlama AI About Your Timetable")
@@ -107,37 +104,31 @@ st.markdown("## 📅 Auto-Schedule Missing Timetable Slots")
 selected_file = st.selectbox("Choose a timetable file to auto-fill missing slots:", list(uploaded_files.keys()))
 
 if st.button("Auto-Schedule"):
-    result = auto_schedule(uploaded_files[selected_file])
-    if isinstance(result, Document):  # DOCX file
-        output = BytesIO()
-        result.save(output)
-        output.seek(0)
-        st.download_button("📥 Download Updated DOCX", output, file_name="updated_timetable.docx", mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
+    if uploaded_files[selected_file]:
+        scheduled_file = auto_schedule(uploaded_files[selected_file])
+        st.download_button("📥 Download Auto-Scheduled Timetable", scheduled_file, file_name=f"{selected_file}_auto.xlsx", mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
     else:
-        st.write("✅", result)
+        st.write("⚠️ Please upload a file first.")
 
-# ✅ Display Uploaded Timetables
+# ✅ Display & Edit Uploaded Timetables (DOCX + XLSX)
 st.markdown("## 📜 View & Edit Uploaded Timetables")
 
 for name, file in uploaded_files.items():
     if file:
         if file.name.endswith(".xlsx"):
-            df = pd.read_excel(file)
+            df = load_xlsx(file)
             edited_df = st.data_editor(df, num_rows="dynamic")
-            st.download_button("📥 Download Edited Excel", edited_df.to_csv(index=False).encode(), file_name=f"{name}.csv", mime="text/csv")
+            csv_output = edited_df.to_csv(index=False).encode()
+            st.download_button("📥 Download Edited Excel", csv_output, file_name=f"{name}.csv", mime="text/csv")
+
         elif file.name.endswith(".docx"):
-            doc = Document(file)
-            doc_text = "\n".join([para.text for para in doc.paragraphs])
-            edited_text = st.text_area(f"Edit {name}:", doc_text)
-            new_doc = Document()
-            for line in edited_text.split("\n"):
-                new_doc.add_paragraph(line)
-            output = BytesIO()
-            new_doc.save(output)
-            output.seek(0)
-            st.download_button("📥 Download Edited DOCX", output, file_name=f"{name}.docx", mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
-
-# ✅ PDF Export Feature
+            doc_text = load_docx(file)
+            edited_text = st.text_area(f"📄 Edit {name}:", doc_text, height=400)
+            if st.button(f"💾 Save & Download {name}"):
+                docx_output = save_docx(edited_text)
+                st.download_button("📥 Download Edited DOCX", docx_output, file_name=f"{name}.docx", mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
+
+# ✅ Export to PDF
 st.markdown("## 📄 Export to PDF")
 if st.button("Export to PDF"):
     pdf_output = BytesIO()
@@ -151,4 +142,4 @@ if st.button("Export to PDF"):
     c.save()
     pdf_output.seek(0)
     st.download_button("📥 Download PDF", pdf_output, file_name="timetable.pdf", mime="application/pdf")
-
+