from transformers import pipeline
import gradio as gr

pipe = pipeline(task="text-generation", model="LiquidAI/LFM2.5-350M", dtype="auto", device_map="auto")

starter = {"role": "system", "content": "You are a helpful general purpose chatbot. Answer user's questions as accurately and concisely as possible."}

def chat_with_model(message, history):
  if not history:
    history.append(starter)
  conversation = history.copy()
  print(f"History: {history}")
  print(f"User Message: {message}")

  conversation.append({"role": "user", "content": message})

  response = pipe(conversation, max_new_tokens=100)

  return response[0]["generated_text"][-1]["content"]

demo = gr.ChatInterface(
  fn=chat_with_model,
  title="CAP 6640 NLP Assignment 4", 
  description="Mini Chatbot using HF",
  examples=["Explain Gravity", "Why is the sky blue?"]
)

demo.launch()