Update app.py

app.py

@@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain.prompts import PromptTemplate
 from langchain.llms import HuggingFacePipeline
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 import torch
 import os
 import requests
@@ -53,16 +53,19 @@ def create_vector_store(transcript):
     return FAISS.from_documents(docs, embeddings)
 
 
-# 🤖 Load Free BLOOM locally
+# 🤖 Load Free Flan-T5 locally (Better than BLOOM)
 @st.cache_resource
-def load_bloom():
-    model_name = "bigscience/bloom-560m"
+def load_flan_t5():
+    model_name = "google/flan-t5-base"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     pipe = pipeline(
-        "text-generation",
+        "text2text-generation",
         model=model,
         tokenizer=tokenizer,
+        max_length=512,
+        temperature=0.7,
+        do_sample=True,
         device=0 if torch.cuda.is_available() else -1
     )
     return HuggingFacePipeline(pipeline=pipe)
@@ -72,26 +75,33 @@ def load_bloom():
 def build_model(model_choice, temperature=0.7):
     if model_choice == "DeepSeek":
         repo_id = "deepseek-ai/DeepSeek-V3.2-Exp"  # paid
-        llm = HuggingFaceEndpoint(repo_id=repo_id, huggingfacehub_api_token=api_key, task="text-generation")
+        llm = HuggingFaceEndpoint(
+            repo_id=repo_id, 
+            huggingfacehub_api_token=api_key, 
+            task="text-generation",
+            max_new_tokens=500
+        )
         return ChatHuggingFace(llm=llm, temperature=temperature)
     elif model_choice == "OpenAI":
         repo_id = "openai/gpt-oss-20b"  # paid
-        llm = HuggingFaceEndpoint(repo_id=repo_id, huggingfacehub_api_token=api_key, task="text-generation")
+        llm = HuggingFaceEndpoint(
+            repo_id=repo_id, 
+            huggingfacehub_api_token=api_key, 
+            task="text-generation",
+            max_new_tokens=500
+        )
         return ChatHuggingFace(llm=llm, temperature=temperature)
     else:
-        return load_bloom()  # free local BLOOM
+        return load_flan_t5()  # free local Flan-T5
 
 
-# 🧾 Prompt Template
+# 🧾 Prompt Template (Simplified for T5)
 prompt_template = PromptTemplate(
     template=(
-        "You are a helpful assistant.\n\n"
-        "Answer the question using the context provided below.\n"
-        "If the context does not mention the topic, say clearly: 'There is no mention of the topic in the video you provided.'\n"
-        "Then, based on your own knowledge, try to answer the question.\n"
-        "If both the context and your knowledge are insufficient, say: 'I don't know.'\n\n"
-        "Context:\n{context}\n\n"
-        "Question:\n{question}"
+        "Answer the question based on the context below.\n\n"
+        "Context: {context}\n\n"
+        "Question: {question}\n\n"
+        "Answer:"
     ),
     input_variables=["context", "question"]
 )
@@ -102,7 +112,7 @@ st.title("🎥 YouTube Transcript Chatbot (Hybrid: Free + Paid)")
 
 video_id = st.text_input("YouTube Video ID", value="lv1_-RER4_I")
 query = st.text_area("Your Query", value="What is RAG?")
-model_choice = st.radio("Model to Use", ["DeepSeek", "OpenAI", "Free BLOOM"])
+model_choice = st.radio("Model to Use", ["Flan-T5 (Free)", "DeepSeek", "OpenAI"])
 temperature = st.slider("Temperature", 0, 100, value=50) / 100.0
 
 if st.button("🚀 Run Chatbot"):
@@ -117,11 +127,21 @@ if st.button("🚀 Run Chatbot"):
                 st.success(f"✅ Transcript fetched! ({len(transcript)} characters)")
 
                 with st.spinner("Generating response..."):
-                    retriever = create_vector_store(transcript).as_retriever(search_type="mmr", search_kwargs={"k": 5})
+                    retriever = create_vector_store(transcript).as_retriever(
+                        search_type="mmr", 
+                        search_kwargs={"k": 5}
+                    )
                     relevant_docs = retriever.invoke(query)
                     context_text = "\n\n".join(doc.page_content for doc in relevant_docs)
                     prompt = prompt_template.format(context=context_text, question=query)
 
                     model = build_model(model_choice, temperature)
-                    response = model.invoke(prompt) if model_choice != "Free BLOOM" else model(prompt)
-                    st.text_area("Model Response", value=response if isinstance(response, str) else response.content, height=400)
+                    
+                    # Handle response based on model type
+                    if model_choice == "Flan-T5 (Free)":
+                        response = model(prompt)
+                    else:
+                        response_obj = model.invoke(prompt)
+                        response = response_obj.content if hasattr(response_obj, 'content') else str(response_obj)
+                    
+                    st.text_area("Model Response", value=response, height=400)