Update app.py

app.py

@@ -1,58 +1,49 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
-import torch
 from transformers import AutoTokenizer, pipeline
-from peft import AutoPeftModelForCausalLM  # comes from `peft`
-
+from peft import AutoPeftModelForCausalLM
 
+# HF repo of your LoRA-finetuned model (the one AutoTrain pushed)
+FINETUNED_MODEL_ID = "janajankovic/autotrain-juhh6-uwiv9"  # <<< CHANGE THIS TO YOUR REPO ID
 
-BASE_MODEL = "cjvt/GaMS-1B-Chat"
-ADAPTER_ID = "janajankovic/autotrain-juhh6-uwiv9"
 
-# Tokenizer from the base model
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+# Load base+LoRA via PEFT
+model = AutoPeftModelForCausalLM.from_pretrained(FINETUNED_MODEL_ID)
+base_model_id = model.config.base_model_name_or_path
 
-# Load base+adapter via AutoPeftModel
-model = AutoPeftModelForCausalLM.from_pretrained(
-    ADAPTER_ID,
-    torch_dtype=torch.float32,  # CPU in this Space
-)
+# Use tokenizer from the base model (GaMS-1B-Chat)
+tokenizer = AutoTokenizer.from_pretrained(base_model_id)
 
-pipe = pipeline(
+# Text generation pipeline
+text_gen = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
     max_new_tokens=256,
     do_sample=True,
-    top_p=0.9,
     temperature=0.7,
+    top_p=0.9,
 )
 
 
-def chat_fn(history, user_input):
-    # simple prompt: previous turns concatenated
-    text = ""
-    for turn in history:
-        text += f"Uporabnik: {turn[0]}\nModel: {turn[1]}\n"
-    text += f"Uporabnik: {user_input}\nModel:"
-
-    out = pipe(text)[0]["generated_text"]
-    reply = out.split("Model:")[-1].strip()
-    history.append((user_input, reply))
-    return history, ""
+def respond(message, history):
+    # message: current user message (string)
+    # history: list of [user, assistant] pairs (ignored here, minimal chat)
+    prompt = message
+    outputs = text_gen(prompt, num_return_sequences=1)
+    text = outputs[0]["generated_text"]
 
+    # Many causal LM heads echo the prompt; strip it out if present
+    if text.startswith(prompt):
+        text = text[len(prompt):].lstrip()
 
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(label="Slovenski finetune")
-    msg = gr.Textbox(label="Vnos")
-    clear = gr.Button("Počisti")
+    # ChatInterface expects a plain string here
+    return text
 
-    def user_send(message, chat_history):
-        chat_history = chat_history or []
-        return "", chat_fn(chat_history, message)[0]
 
-    msg.submit(user_send, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: ([], ""), None, [chatbot, msg], queue=False)
+demo = gr.ChatInterface(
+    fn=respond,
+    title="GenUI – Slovene fine-tuned chat",
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()