Update app.py

app.py

@@ -1,51 +1,45 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load model
-tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
-model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
+# Load TinyLlama (small but smart)
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 def generate_notes(topic):
-    """
-    Generate detailed paragraph-style study notes for a given topic.
-    """
-    if not topic.strip():
+    topic = topic.strip()
+    if not topic:
         return "⚠️ Please enter a topic first."
 
     prompt = (
-        f"Write detailed, well-structured study notes about the topic '{topic}'. "
-        f"Include an introduction, key concepts, examples, and a conclusion. "
-        f"Make the notes clear, paragraph-based, and easy for a student to understand. "
-        f"Use complete sentences and proper formatting."
+        f"You are an expert teacher. Write detailed, well-structured study notes about '{topic}'. "
+        f"Include an introduction, main concepts, key points, and a short summary. "
+        f"Keep the tone educational and clear for students."
     )
 
-    # Encode input
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-
-    # Generate long, detailed text
+    inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(
         **inputs,
-        max_new_tokens=600,       # increase for longer output
-        temperature=0.8,          # adds variety
-        top_p=0.9,                # nucleus sampling
-        repetition_penalty=1.1,   # prevents repeating phrases
-        do_sample=True            # enables more creative generation
+        max_new_tokens=600,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.2,
     )
 
-    notes = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return notes.strip()
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return text.replace(prompt, "").strip()
 
-# Gradio UI
 iface = gr.Interface(
     fn=generate_notes,
-    inputs=gr.Textbox(lines=2, placeholder="Enter a topic (e.g. Photosynthesis, Machine Learning)"),
-    outputs=gr.Textbox(label="Generated Detailed Notes", lines=15),
-    title="📘 AI Detailed Note Generator",
-    description="Generate detailed, structured notes with multiple paragraphs using FLAN-T5-large (runs free on CPU).",
+    inputs=gr.Textbox(lines=2, placeholder="Enter a topic (e.g. Photosynthesis, World Geography)"),
+    outputs=gr.Textbox(lines=15, label="Generated Notes"),
+    title="📚 AI Note Generator (TinyLlama 1.1B Chat)",
+    description="Generate long, educational notes for free using the TinyLlama 1.1B Chat model.",
     examples=[
-        ["Quantum computing"],
-        ["Climate change and its effects"],
-        ["Photosynthesis process"]
+        ["Photosynthesis process"],
+        ["Causes of World War II"],
+        ["Artificial Intelligence"],
+        ["The Water Cycle"]
     ]
 )