Update app.py

app.py

@@ -1,48 +1,174 @@
+import os
+
 import gradio as gr
-from transformers import AutoTokenizer, pipeline
-from peft import AutoPeftModelForCausalLM
+import pandas as pd
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from peft import PeftModel
+
+# -------------------------------------------------------------------
+# CONFIG
+# -------------------------------------------------------------------
+
+# Your fine-tuned adapter repo on HF
+MODEL_ID = "janajankovic/autotrain-juhh6-uwiv9"  # change if needed
 
-# 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 that was fine-tuned (the one you used in AutoTrain)
+BASE_MODEL_ID = "cjvt/GaMS-1B-Chat"  # change if different
 
+# CSV with chunks (already in the Space repo)
+CSV_PATH = "chunks_for_autotrain.csv"
 
-# Load base+LoRA via PEFT
-model = AutoPeftModelForCausalLM.from_pretrained(FINETUNED_MODEL_ID)
-base_model_id = model.config.base_model_name_or_path
+# How many *extra* chunks (besides the top-1) to add
+N_NEIGHBORS = 4
 
-# Use tokenizer from the base model (GaMS-1B-Chat)
-tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+MAX_NEW_TOKENS = 256
+TEMPERATURE = 0.7
+TOP_P = 0.9
+
+# -------------------------------------------------------------------
+# LOAD MODEL (BASE + PEFT ADAPTER)
+# -------------------------------------------------------------------
+
+print("Loading base model and tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_ID,
+    torch_dtype="auto",
+)
 
-# Text generation pipeline
-text_gen = pipeline(
+# Attach LoRA / PEFT adapter
+print("Loading PEFT adapter...")
+model = PeftModel.from_pretrained(base_model, MODEL_ID)
+
+# Make sure pad token is set
+if model.config.pad_token_id is None and model.config.eos_token_id is not None:
+    model.config.pad_token_id = model.config.eos_token_id
+
+generator = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_new_tokens=256,
-    do_sample=True,
-    temperature=0.7,
-    top_p=0.9,
 )
 
+# -------------------------------------------------------------------
+# LOAD CHUNKS + BUILD TF-IDF RETRIEVER
+# -------------------------------------------------------------------
+
+print("Loading CSV chunks...")
+df = pd.read_csv(CSV_PATH)
+df["text"] = df["text"].fillna("")
+
+documents = df["text"].tolist()
+
+print("Building TF-IDF index...")
+vectorizer = TfidfVectorizer(max_features=50000)
+doc_matrix = vectorizer.fit_transform(documents)
+
+# -------------------------------------------------------------------
+# RETRIEVAL: TOP-1 + NEXT N_NEIGHBORS MOST SIMILAR CHUNKS
+# -------------------------------------------------------------------
+
+def retrieve_chunks(query: str, n_neighbors: int = N_NEIGHBORS):
+    query = query.strip()
+    if not query:
+        return []
+
+    # similarity of question vs all chunks
+    q_vec = vectorizer.transform([query])
+    sims = cosine_similarity(q_vec, doc_matrix).flatten()
+
+    if sims.max() <= 0:
+        return []
+
+    # indices sorted by similarity to the question (desc)
+    sorted_indices = sims.argsort()[::-1]
+
+    # central: most similar to question
+    central_idx = int(sorted_indices[0])
+
+    # neighbors: next n_neighbors most similar to question
+    neighbor_indices = [central_idx]
+    for idx in sorted_indices[1:]:
+        if len(neighbor_indices) >= n_neighbors + 1:
+            break
+        neighbor_indices.append(int(idx))
+
+    # keep order: central first, then neighbors
+    selected_texts = [documents[i] for i in neighbor_indices]
+    return selected_texts
+
+def build_context(question: str) -> str:
+    chunks = retrieve_chunks(question, N_NEIGHBORS)
+    if not chunks:
+        return ""
+
+    # Optional: prefix chunks for clarity (not strictly needed)
+    labelled = []
+    for i, ch in enumerate(chunks):
+        labelled.append(f"[CHUNK {i+1}]\n{ch}")
+    return "\n\n".join(labelled)
+
+# -------------------------------------------------------------------
+# CHAT FUNCTION
+# -------------------------------------------------------------------
+
+SYSTEM_PROMPT = (
+    "Ti si pomočnik, ki odgovarja v slovenščini.\n"
+    "Uporabi spodnji kontekst, če je relevanten. "
+    "Če kontekst ne vsebuje odgovora, odgovori po svojih najboljših močeh "
+    "in jasno povej, da se opiraš na splošno znanje.\n"
+)
+
+def generate_answer(message: str) -> str:
+    context = build_context(message)
+
+    if context:
+        full_prompt = (
+            f"{SYSTEM_PROMPT}\n"
+            f"Kontekst:\n{context}\n\n"
+            f"Vprašanje uporabnika:\n{message}\n\n"
+            f"Odgovor (v slovenščini):\n"
+        )
+    else:
+        full_prompt = (
+            f"{SYSTEM_PROMPT}\n"
+            f"Vprašanje uporabnika:\n{message}\n\n"
+            f"Odgovor (v slovenščini):\n"
+        )
+
+    outputs = generator(
+        full_prompt,
+        max_new_tokens=MAX_NEW_TOKENS,
+        do_sample=True,
+        temperature=TEMPERATURE,
+        top_p=TOP_P,
+        pad_token_id=model.config.pad_token_id,
+    )
 
-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"]
+    generated = 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()
+    # strip the prompt from the beginning
+    answer = generated[len(full_prompt):].strip()
+    return answer
 
-    # ChatInterface expects a plain string here
-    return text
+def chat_fn(message, history):
+    return generate_answer(message)
 
+# -------------------------------------------------------------------
+# GRADIO UI
+# -------------------------------------------------------------------
 
 demo = gr.ChatInterface(
-    fn=respond,
-    title="GenUI – Slovene fine-tuned chat",
+    fn=chat_fn,
+    title="Gen-UI fine-tuned Slovene model",
+    description=(
+        "Klepet z lastnim fine-tunanim modelom.\n"
+        "Model samodejno poišče najbližje besedilne 'chunke' v CSV in jih uporabi kot kontekst."
+    ),
 )
 
 if __name__ == "__main__":