Update app.py

app.py

@@ -1,47 +1,29 @@
 import gradio as gr
-from langchain_community.chat_models import HuggingFaceChat
-from langchain.chains import ConversationChain
-from langchain.memory import ConversationBufferMemory
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from langchain_community.llms import HuggingFacePipeline
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
 
-MODEL_ID = "google/gemma-2b"  # Adjust if you have a local path or HF repo
+# Load model and tokenizer (Gemma 2B or similar)
+model_id = "google/gemma-2b-it"
 
-def load_model():
-    # Load Gemma 2B chat model on CPU
-    return HuggingFaceChat(
-        model=MODEL_ID,
-        model_kwargs={"device_map": "cpu"},
-        temperature=0.7,
-    )
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
 
-def main():
-    llm = load_model()
-    memory = ConversationBufferMemory(return_messages=True)
-    conversation = ConversationChain(llm=llm, memory=memory)
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
 
-    def respond(user_input):
-        response = conversation.run(user_input)
-        return response
+llm = HuggingFacePipeline(pipeline=pipe)
 
-    title = "🦚 Meet Krish"
-    description = "A wise, witty, and compassionate friend — Chatbot named KrishWay"
+# Simple prompt template
+prompt = PromptTemplate.from_template("You are Krish, a wise and witty friend.\n\nUser: {question}\nKrish:")
 
-    with gr.Blocks() as demo:
-        gr.Markdown(f"# {title}")
-        gr.Markdown(description)
+chain = LLMChain(prompt=prompt, llm=llm)
 
-        chatbot = gr.Chatbot()
-        msg = gr.Textbox(placeholder="Say something to Krish...")
-        clear = gr.Button("Clear")
+# Gradio interface
+def chat_fn(message):
+    response = chain.run({"question": message})
+    return response.strip()
 
-        def user_message(text, chat_history):
-            bot_response = respond(text)
-            chat_history = chat_history + [(text, bot_response)]
-            return "", chat_history
+iface = gr.Interface(fn=chat_fn, inputs="text", outputs="text", title="🦚 Meet Krish", description="A wise, witty, and compassionate friend - KrishWay")
 
-        msg.submit(user_message, [msg, chatbot], [msg, chatbot])
-        clear.click(lambda: [], None, chatbot)
-
-    demo.launch()
-
-if __name__ == "__main__":
-    main()
+iface.launch()