Update app.py

app.py

@@ -1,11 +1,15 @@
 import json
 import requests
-import openpyxl
-import pandas as pd
+import os
 import streamlit as st
-from io import BytesIO
+import pandas as pd
+import openpyxl
+import torch
 from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
+from transformers import pipeline
+from docx import Document
+from io import BytesIO
 
 # ✅ Streamlit UI Setup
 st.set_page_config(page_title="AI-Powered Timetable", layout="wide")
@@ -15,31 +19,35 @@ st.markdown("<h1 style='text-align: center; color: #4CAF50;'>📅 AI-Powered Tim
 st.sidebar.markdown("## 🔑 Enter Hugging Face API Key")
 hf_api_key = st.sidebar.text_input("API Key", type="password")
 
-# ✅ File Upload Section
+# ✅ File Upload Section (Excel & DOCX)
 st.sidebar.markdown("## 📂 Upload Your Timetable Files")
-uploaded_file = st.sidebar.file_uploader("Upload Timetable", type=["xlsx"])
-
-# ✅ Load Timetable Data
+uploaded_files = {
+    "Master Timetable": st.sidebar.file_uploader("Upload Master Timetable", type=["xlsx", "docx"]),
+    "Lab Timetable": st.sidebar.file_uploader("Upload Lab Timetable", type=["xlsx", "docx"]),
+    "Classroom Timetable": st.sidebar.file_uploader("Upload Classroom Timetable", type=["xlsx", "docx"]),
+    "Individual Timetable": st.sidebar.file_uploader("Upload Individual Timetable", type=["xlsx", "docx"]),
+}
+
+# ✅ Load Timetable Data (Excel & DOCX)
 def load_timetable(file):
-    if file:
-        df = pd.read_excel(file)
-        return df
-    return None
-
-df = load_timetable(uploaded_file)
-
-# ✅ Ask TinyLlama API
+    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]
+
+# ✅ Ask TinyLlama AI via API
 def ask_tinyllama_api(query):
     if not hf_api_key:
         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
-    }
+    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:
@@ -47,65 +55,100 @@ def ask_tinyllama_api(query):
     else:
         return f"API Error: {response.status_code} - {response.text}"
 
-# ✅ Auto-Schedule Missing Slots
-def auto_schedule(df):
-    if df is None:
+# ✅ Auto-Schedule Missing Slots (Excel & DOCX)
+def auto_schedule(file):
+    if not file:
         return "No timetable uploaded."
-    
-    empty_slots = df[df.isnull().any(axis=1)].index.tolist()
-    for idx in empty_slots:
-        query = f"Suggest a subject and faculty for the empty slot in row {idx+2}."
-        suggestion = ask_tinyllama_api(query)
-        try:
-            subject, faculty = suggestion.split(", Faculty: ")
-            df.at[idx, "Subject"] = subject.strip()
-            df.at[idx, "Faculty"] = faculty.strip()
-        except:
-            continue
-    
-    return f"Auto-scheduling completed for {len(empty_slots)} slots."
-
-# ✅ Convert DataFrame to PDF
-def convert_to_pdf(df):
-    buffer = BytesIO()
-    pdf = canvas.Canvas(buffer, pagesize=letter)
-    pdf.setFont("Helvetica", 10)
-
-    pdf.drawString(30, 750, "AI-Powered Timetable")
 
-    y_position = 730
-    for i, row in df.iterrows():
-        pdf.drawString(30, y_position, str(row.values))
-        y_position -= 20
-        if y_position < 50:
-            pdf.showPage()
-            pdf.setFont("Helvetica", 10)
-            y_position = 750
-    
-    pdf.save()
-    buffer.seek(0)
-    return buffer
-
-# ✅ Display and Edit Timetable
-if df is not None:
-    st.markdown("## 📝 Edit Your Timetable")
-    edited_df = st.data_editor(df, key="editable_table")
-
-    # ✅ Auto-Schedule Button
-    if st.button("Auto-Schedule Missing Slots"):
-        result = auto_schedule(edited_df)
-        st.write("✅", result)
-
-    # ✅ Download as PDF
-    if st.button("Download as PDF"):
-        pdf_buffer = convert_to_pdf(edited_df)
-        st.download_button("📥 Download Timetable PDF", pdf_buffer, file_name="timetable.pdf", mime="application/pdf")
+    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
 
 # ✅ AI Query Section
-st.markdown("## 🤖 Ask TinyLlama About Your Timetable")
+st.markdown("## 🤖 Ask TinyLlama AI About Your Timetable")
 user_query = st.text_input("Type your question here (e.g., 'Who is free at 10 AM on Monday?')")
 
 if st.button("Ask AI via API"):
     ai_response = ask_tinyllama_api(user_query)
     st.write("🧠 **TinyLlama Suggests:**", ai_response)
-    
+
+# ✅ Auto-Schedule Feature
+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")
+    else:
+        st.write("✅", result)
+
+# ✅ Display Uploaded Timetables
+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)
+            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")
+        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
+st.markdown("## 📄 Export to PDF")
+if st.button("Export to PDF"):
+    pdf_output = BytesIO()
+    c = canvas.Canvas(pdf_output, pagesize=letter)
+    c.drawString(100, 750, "AI-Powered Timetable Export")
+    y_position = 730
+    for name, file in uploaded_files.items():
+        if file:
+            c.drawString(100, y_position, f"{name}: {file.name}")
+            y_position -= 20
+    c.save()
+    pdf_output.seek(0)
+    st.download_button("📥 Download PDF", pdf_output, file_name="timetable.pdf", mime="application/pdf")
+