add local model

app.py

@@ -6,6 +6,7 @@ from sentence_transformers import SentenceTransformer
 from openai import OpenAI
 import random
 import prompts
+from utils import get_bnb_config, load_base_model, load_fine_tuned_model, generate_response
 
 st.set_page_config(page_title="AI University")
 
@@ -70,6 +71,10 @@ def fixed_knn_retrieval(question_embedding, context_embeddings, top_k=5, min_k=1
 def sec_to_time(start_time):
     return f"{start_time // 60:02}:{start_time % 60:02}"
 
+
+
+
+
 st.markdown("""
     <style>
     .video-wrapper {
@@ -161,22 +166,29 @@ with st.sidebar:
         # latex_overlap_tokens = latex_chunk_tokens // 4
         latex_overlap_tokens = 0
 
-    st.write(' ')
-    with st.expander('Expert model',expanded=False):
-        # st.write('**Expert model**')
-        # with st.container(border=True):
-        # Choose the LLM model
+    st.write(' ')      
+    with st.expander('Expert model', expanded=False):
 
-        use_expert_answer = st.toggle("Use expert answer", value=True)
-        show_expert_responce = st.toggle("Show initial expert answer", value=False)
+            use_expert_answer = st.toggle("Use expert answer", value=True)
+            show_expert_responce = st.toggle("Show initial expert answer", value=False)
 
-        model = st.selectbox("Choose the LLM model", ["gpt-4o-mini", "gpt-3.5-turbo"], key='a1model')
+            model = st.selectbox("Choose the LLM model", ["gpt-4o-mini", "gpt-3.5-turbo", "llama-tommi-0.35"], key='a1model')
 
-        # Temperature
-        expert_temperature = st.slider("Temperature", 0.0, 0.3, .2,  help="Defines the randomness in the next token prediction. Lower: More predictable and focused. Higher: More adventurous and diverse.", key='a1t')
+            if model == "llama-tommi-0.35":
+                tommi_do_sample = st.toggle("Enable Sampling", value=True, key='tommi_sample')
 
-        expert_top_p = st.slider("Top P", 0.1, 0.3, 0.1, help="Defines the range of token choices the model can consider in the next prediction. Lower: More focused and restricted to high-probability options. Higher: More creative, allowing consideration of less likely options.", key='a1p')
-        
+                if tommi_do_sample:
+                    tommi_temperature = st.slider("Temperature", 0.0, 1.5, 0.7, key='tommi_temp')
+                    tommi_top_k = st.slider("Top K", 0, 100, 50, key='tommi_top_k')
+                    tommi_top_p = st.slider("Top P", 0.0, 1.0, 0.95, key='tommi_top_p')
+                else:
+                    tommi_num_beams = st.slider("Num Beams", 1, 10, 4, key='tommi_num_beams')
+
+                tommi_max_new_tokens = st.slider("Max New Tokens", 100, 2000, 500, step=50, key='tommi_max_new_tokens')
+            else:
+                expert_temperature = st.slider("Temperature", 0.0, 1.5, 0.7, key='a1t')
+                expert_top_p = st.slider("Top P", 0.0, 1.0, 0.9, key='a1p')
+                expert_top_k = st.slider("Top K", 0, 100, 50, key='a1k')
 
     with st.expander('Synthesis model',expanded=False):
 
@@ -281,9 +293,41 @@ if submit_button_placeholder.button("AI Answer", type="primary"):
                     context += context_item['text'] + '\n\n'
             
             if use_expert_answer:
-                st.session_state.expert_answer = prompts.openai_domain_specific_answer_generation("Finite Element Method", st.session_state.question, model=model, temperature=expert_temperature, top_p=expert_top_p)
+                if model == "llama-tommi-0.35":
+                    if 'tommi_model' not in st.session_state:
+                        tommi_model, tommi_tokenizer = load_fine_tuned_model(adapter_path, base_model_path)
+                        st.session_state.tommi_model = tommi_model
+                        st.session_state.tommi_tokenizer = tommi_tokenizer
+
+                    messages = [
+                        {"role": "system", "content": "You are an expert in Finite Element Methods."},
+                        {"role": "user", "content": st.session_state.question}
+                    ]
+
+                    st.session_state.expert_answer = generate_response(
+                        model=st.session_state.tommi_model,
+                        tokenizer=st.session_state.tommi_tokenizer,
+                        messages=messages,
+                        do_sample=tommi_do_sample,
+                        temperature=tommi_temperature if tommi_do_sample else None,
+                        top_k=tommi_top_k if tommi_do_sample else None,
+                        top_p=tommi_top_p if tommi_do_sample else None,
+                        num_beams=tommi_num_beams if not tommi_do_sample else 1,
+                        max_new_tokens=tommi_max_new_tokens
+                    )
+                else:
+                    st.session_state.expert_answer = prompts.openai_domain_specific_answer_generation(
+                        "Finite Element Method",
+                        st.session_state.question,
+                        model=model,
+                        temperature=expert_temperature,
+                        top_p=expert_top_p,
+                        top_k=expert_top_k
+                    )
             else:
                 st.session_state.expert_answer = 'No Expert Answer. Only use the context.'
+
+
             answer = prompts.openai_context_integration("Finite Element Method", st.session_state.question, st.session_state.expert_answer, context, model=model, temperature=integration_temperature, top_p=integration_top_p)
 
         if answer.split()[0] == "NOT_ENOUGH_INFO":

utils.py

@@ -0,0 +1,141 @@
+import torch
+from transformers import BitsAndBytesConfig, AutoModelForCausalLM, PreTrainedTokenizerFast
+from peft import PeftModel
+
+#-----------------------------------------
+# Quantization Config
+#-----------------------------------------
+def get_bnb_config():
+    return BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_storage=torch.float16
+    )
+
+#-----------------------------------------
+# Base Model Loader
+#-----------------------------------------
+def load_base_model(base_model_path: str):
+    """
+    Loads a base LLM model with 4-bit quantization and tokenizer.
+    
+    Args:
+        base_model_path (str): HF model path
+    
+    Returns:
+        model (AutoModelForCausalLM)
+        tokenizer (PreTrainedTokenizerFast)
+    """
+    bnb_config = get_bnb_config()
+
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(base_model_path, return_tensors="pt")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_path,
+        quantization_config=bnb_config,
+        trust_remote_code=True,
+        attn_implementation="eager",
+        torch_dtype=torch.float16
+    )
+
+    return model, tokenizer
+
+#-----------------------------------------
+# Fine-Tuned Model Loader
+#-----------------------------------------
+def load_fine_tuned_model(adapter_path: str, base_model_path: str):
+    """
+    Loads the fine-tuned model by applying LoRA adapter to a base model.
+    
+    Args:
+        adapter_path (str): Local or HF adapter path
+        base_model_path (str): Base LLM model path
+    
+    Returns:
+        fine_tuned_model (PeftModel)
+        tokenizer (PreTrainedTokenizerFast)
+    """
+    bnb_config = get_bnb_config()
+
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(base_model_path, return_tensors="pt")
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_model_path,
+        quantization_config=bnb_config,
+        trust_remote_code=True,
+        attn_implementation="eager",
+        torch_dtype=torch.float16
+    )
+
+    fine_tuned_model = PeftModel.from_pretrained(
+        base_model,
+        adapter_path,
+        device_map="auto"
+    )
+
+    return fine_tuned_model, tokenizer
+
+#-----------------------------------------
+# Inference Function
+#-----------------------------------------
+@torch.no_grad()
+def generate_response(
+    model: AutoModelForCausalLM,
+    tokenizer: PreTrainedTokenizerFast,
+    messages: list,
+    do_sample: bool = False,
+    temperature: float = 0.7,
+    top_k: int = 50,
+    top_p: float = 0.95,
+    num_beams: int = 1,
+    max_new_tokens: int = 500
+) -> str:
+    """
+    Runs inference on an LLM model.
+
+    Args:
+        model (AutoModelForCausalLM)
+        tokenizer (PreTrainedTokenizerFast)
+        messages (list): List of dicts containing 'role' and 'content'
+    
+    Returns:
+        str: Model response
+    """
+    # Ensure pad token exists
+    tokenizer.pad_token = "<|reserved_special_token_5|>"
+
+    # Create chat prompt
+    input_text = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=False
+    )
+
+    # Tokenize input
+    inputs = tokenizer(
+        input_text,
+        max_length=500,
+        truncation=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    generation_params = {
+        "do_sample": do_sample,
+        "temperature": temperature if do_sample else None,
+        "top_k": top_k if do_sample else None,
+        "top_p": top_p if do_sample else None,
+        "num_beams": num_beams if not do_sample else 1,
+        "max_new_tokens": max_new_tokens
+    }
+
+    output = model.generate(**inputs, **generation_params)
+
+    # Decode and clean up response
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    if 'assistant' in response:
+        response = response.split('assistant')[1].strip()
+
+    return response