Update app.py

app.py

@@ -1,35 +1,65 @@
 import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
-# 1) Load your tokenizer
-tokenizer = AutoTokenizer.from_pretrained("testi123456789/elektromart")
+# 1) Load tokenizer and base model on CPU (or GPU if available)
+tokenizer = AutoTokenizer.from_pretrained("finnish-nlp/ahma-3b")
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
 
-# 2) Load the base model onto CPU (full-precision)
-base = AutoModelForCausalLM.from_pretrained("finnish-nlp/ahma-3b")
+base_model = AutoModelForCausalLM.from_pretrained(
+    "finnish-nlp/ahma-3b",
+    torch_dtype=torch.float32,
+    device_map={"": "cpu"}
+)
 
-# 3) Apply your LoRA adapter without any device dispatching
+# 2) Apply your fine-tuned LoRA adapter
 model = PeftModel.from_pretrained(
-    base,
-    "testi123456789/elektromart",
-    device_map=None
+    base_model,
+    "testi123456789/elektromart"
 )
 model.to("cpu")
 model.eval()
 
-# 4) Define the chat function
-def chat_fn(prompt):
+# 3) Instruction you fine-tuned on
+INSTRUCTION = "Vastaa asiakkaan kyselyyn ystävällisesti ElektroMartin asiakaspalveluna."
+
+def chat_fn(user_question: str, max_new_tokens: int = 100,
+            temperature: float = 0.7, repetition_penalty: float = 1.25) -> str:
+    # 4) Build the prompt exactly as during training
+    prompt = f"[INST] {INSTRUCTION}\n{user_question} [/INST]\n"
+
+    # 5) Tokenize & clean up
     inputs = tokenizer(prompt, return_tensors="pt")
     inputs.pop("token_type_ids", None)
     inputs = {k: v.to("cpu") for k, v in inputs.items()}
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+    # 6) Generate
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            pad_token_id=tokenizer.eos_token_id,
+            do_sample=True,
+            repetition_penalty=repetition_penalty
+        )
 
-# 5) Launch Gradio
-gr.Interface(
+    # 7) Decode only the newly generated part
+    generated = outputs[0][ inputs["input_ids"].shape[-1] : ]
+    answer = tokenizer.decode(generated, skip_special_tokens=True)
+    return answer.strip()
+
+# 8) Expose Gradio interface
+iface = gr.Interface(
     fn=chat_fn,
-    inputs=gr.Textbox(placeholder="Kysy jotain…"),
-    outputs="text",
-    title="ElektroMart Chatbot"
-).launch()
+    inputs=[
+        gr.Textbox(label="Kysy jotain…", placeholder="Kirjoita kysymyksesi tähän"),
+    ],
+    outputs=gr.Textbox(label="Vastaus"),
+    title="ElektroMartin Chatbotti"
+)
+
+if __name__ == "__main__":
+    iface.launch()
+