Update app.py

app.py

@@ -1,59 +1,58 @@
-import gradio as gr
-from transformers import pipeline
-from langchain_community.llms import HuggingFacePipeline
-from langchain_core.prompts import PromptTemplate
+import os
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import LLMChain
+from langchain_core.prompts import PromptTemplate
+from langchain_community.llms import HuggingFacePipeline
+import gradio as gr
+
+# Load tokenizer and model
+model_id = "OpenAssistant/oasst-sft-1-pythia-12b"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
 
-# Load HF pipeline (CPU-only model like flan-t5-base)
 pipe = pipeline(
-    "text2text-generation",
-    model="google/flan-t5-base",
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
     max_new_tokens=150,
     do_sample=True,
     temperature=0.7,
-    top_p=0.9
+    top_k=50,
+    top_p=0.9,
+    repetition_penalty=1.1,
+    device=-1  # Force CPU
 )
 
-llm = HuggingFacePipeline(pipeline=pipe)
+# Define Krish personality
+template = """You are Krish — a warm, witty, and wise friend inspired by Lord Krishna from the Puranas.
+Krish speaks with compassion, intelligence, and playful charm.
+He helps users with any problem or question, offering support like a true friend.
+Use gentle wisdom and occasional ancient Indian metaphors or proverbs.
 
-# Prompt Template
-template = (
-    "You are Krish — a warm, witty, and wise friend inspired by Lord Krishna from the Puranas. "
-    "Respond with compassion, intelligence, and charm. Use metaphors or ancient Indian wisdom sometimes.\n\n"
-    "{chat_history}\n"
-    "User: {user_message}\nKrish:"
-)
-prompt = PromptTemplate(
-    input_variables=["chat_history", "user_message"],
-    template=template
-)
+{chat_history}
+User: {user_message}
+Krish:"""
 
-# Memory for chat history
+prompt = PromptTemplate(input_variables=["chat_history", "user_message"], template=template)
 memory = ConversationBufferMemory(memory_key="chat_history")
 
-# LangChain LLMChain
-llm_chain = LLMChain(
-    llm=llm,
-    prompt=prompt,
-    memory=memory,
-    verbose=False,
-)
+llm = HuggingFacePipeline(pipeline=pipe)
+llm_chain = LLMChain(llm=llm, prompt=prompt, memory=memory, verbose=False)
 
-# Handle chat
+# Define Gradio interface logic
 def chat(user_message, history):
-    memory.chat_memory.clear()  # Always sync history
-    for h in history:
-        memory.chat_memory.add_user_message(h[0])
-        memory.chat_memory.add_ai_message(h[1])
     response = llm_chain.predict(user_message=user_message)
     return response
 
-# Gradio ChatInterface with retry + clear
 chatbot = gr.ChatInterface(
     fn=chat,
-  title="🦚 Meet Krish",
-    description="A wise, witty, and compassionate friend inspired by Lord Krishna. ",
+   title="🦚 Meet Krish",
+    description="A wise, witty, and compassionate friend ",
+    theme="default",
+    textbox=gr.Textbox(placeholder="Say something to Krish...", scale=7),
+    # clear_btn="Clear",
 )
 
 if __name__ == "__main__":