Update app.py

app.py

@@ -1,85 +1,98 @@
-import os
 import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 
-MODEL = "gpt-4o-mini"
+# ---- CONFIG ----
+MODEL_ID = "md-nishat-008/TigerLLM-1B-it"  # Bangla-optimized TigerLLM (instruction tuned)
+DEVICE = 0 if torch.cuda.is_available() else -1
 
-SYSTEM_HI = "तुम एक सहायक हिन्दी सहायक हो। स्पष्ट, संक्षिप्त और सही उत्तर दो।"
-SYSTEM_BN = "তুমি একজন সহায়ক বাংলা সহকারী। সংক্ষিপ্ত, স্পষ্ট এবং সঠিক উত্তর দাও।"
+SYSTEM_BN = "তুমি একজন সহায়ক বাংলা সহকারী। সংক্ষিপ্ত, স্পষ্ট এবং ভদ্রভাবে উত্তর দাও।"
+SYSTEM_HI = "तुम एक सहायक हिन्दी सहायक हो। स्पष्ट और संक्षिप्त उत्तर दो। (मॉडल बंगाली-केंद्रित है, इसलिए हिन्दी गुणवत्ता सीमित हो सकती है।)"
 
-def call_openai(messages):
-    api_key = os.getenv("OPENAI_API_KEY", "").strip()
-    if not api_key:
-        raise RuntimeError("OPENAI_API_KEY is not set in Settings → Variables and secrets")
+pipe = None
 
-    # Lazy import so startup doesn't crash if deps missing
-    from openai import OpenAI
-    client = OpenAI(api_key=api_key)
+def load_pipeline():
+    global pipe
+    if pipe is not None:
+        return pipe
 
-    resp = client.chat.completions.create(
-        model=MODEL,
-        messages=messages,
-        temperature=0.6,
-        max_tokens=400,
+    # Load tokenizer & model (works on CPU Spaces too)
+    tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float32,          # CPU-friendly; HF free Spaces usually run on CPU
+        low_cpu_mem_usage=True,
+        device_map="auto" if DEVICE != -1 else None,
     )
-    return resp.choices[0].message.content
 
-def build_messages(history, lang_label):
-    # lang_label is "हिन्दी" or "বাংলা"
-    lang = "bn" if lang_label == "বাংলা" else "hi"
-    sys = SYSTEM_BN if lang == "bn" else SYSTEM_HI
+    pipe = TextGenerationPipeline(
+        model=model,
+        tokenizer=tok,
+        device=DEVICE
+    )
+    return pipe
 
-    messages = [{"role": "system", "content": sys}]
+def build_prompt(history, lang):
+    # history is list[tuple[str|None, str|None]] of (user, bot)
+    sys = SYSTEM_BN if lang == "bn" else SYSTEM_HI
+    conv = [f"<<SYS>> {sys} </SYS>"]
     for user, bot in history:
         if user:
-            messages.append({"role": "user", "content": user})
+            conv.append(f"<<USER>> {user}")
         if bot:
-            messages.append({"role": "assistant", "content": bot})
-    return messages
+            conv.append(f"<<ASSISTANT>> {bot}")
+    conv.append("<<ASSISTANT>>")  # model should continue here
+    return "\n".join(conv)
+
+def reply(history, lang_label):
+    # lang_label from radio: "বাংলা" or "हिन्दी"
+    lang = "bn" if lang_label == "বাংলা" else "hi"
+
+    generator = load_pipeline()
+    prompt = build_prompt(history, lang)
+
+    out = generator(
+        prompt,
+        max_new_tokens=320,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.05,
+        pad_token_id=generator.tokenizer.eos_token_id,
+        eos_token_id=generator.tokenizer.eos_token_id,
+        num_return_sequences=1,
+    )[0]["generated_text"]
+
+    # Extract only the model's final turn after the last "<<ASSISTANT>>"
+    if "<<ASSISTANT>>" in out:
+        answer = out.split("<<ASSISTANT>>")[-1].strip()
+    else:
+        answer = out.strip()
+
+    # Soft safety: trim over-long outputs
+    answer = answer[:2000]
+    # Tiny post-clean
+    answer = answer.replace("<<USER>>", "").replace("<<SYS>>", "")
+
+    # Put answer into the last tuple
+    last_user, _ = history[-1]
+    history[-1] = (last_user, answer)
+    return history
 
 def on_user_submit(user_text, chat_state):
-    """Append the user message; bot will fill the reply in the next step."""
     if not user_text.strip():
         return "", chat_state
-    chat_state = chat_state + [(user_text, None)]
-    return "", chat_state
-
-def on_bot_reply(chat_state, lang_label):
-    """Call the model and set the reply for the last turn."""
-    try:
-        messages = build_messages(chat_state, lang_label)
-        answer = call_openai(messages)
-        # put the bot answer into the last tuple
-        last_user, _ = chat_state[-1]
-        chat_state[-1] = (last_user, answer)
-        return chat_state
-    except Exception as e:
-        last_user, _ = chat_state[-1]
-        err_hi = f"त्रुट��: {e}"
-        err_bn = f"ত্রুটি: {e}"
-        chat_state[-1] = (last_user, err_bn if lang_label == "বাংলা" else err_hi)
-        return chat_state
+    return "", chat_state + [(user_text, None)]
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("### 🇮🇳 Indic Chat — हिंदी / বাংলা")
-
-    lang = gr.Radio(
-        choices=["हिन्दी", "বাংলা"],
-        value="हिन्दी",
-        label="Language",
-    )
+    gr.Markdown("### 🐯 TigerLLM Chat — বাংলা / हिंदी\n*(Bangla-optimized model; Hindi quality may be limited)*")
 
-    # ✅ Use 'tuples' (plural)
+    lang = gr.Radio(choices=["বাংলা", "हिन्दी"], value="বাংলা", label="Language")
     chat = gr.Chatbot(height=420, type="tuples")
-
-    txt = gr.Textbox(placeholder="Type in Hindi or Bangla…", autofocus=True)
+    txt = gr.Textbox(placeholder="বাংলা বা हिन्दी লিখুন…", autofocus=True)
     send = gr.Button("Send")
 
-    # Wire up: user adds a turn → then bot replies
-    txt.submit(on_user_submit, [txt, chat], [txt, chat]).then(
-        on_bot_reply, [chat, lang], [chat]
-    )
-    send.click(on_user_submit, [txt, chat], [txt, chat]).then(
-        on_bot_reply, [chat, lang], [chat]
-    )
+    txt.submit(on_user_submit, [txt, chat], [txt, chat]).then(reply, [chat, lang], [chat])
+    send.click(on_user_submit, [txt, chat], [txt, chat]).then(reply, [chat, lang], [chat])
 
 demo.launch()