updated app.py

app.py

@@ -1,68 +1,71 @@
 import gradio as gr
-from llama_cpp import Llama
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Load YOUR fine-tuned model
-model_path = "samzito12/lora_model"
+print("Chargement du modèle...")
 
-print("Loading model...")
-llm = Llama.from_pretrained(
-    repo_id=model_path,
-    filename="llama-3.2-3b-instruct.Q8_0.gguf",
-    n_ctx=2048,
-    n_threads=2,
-    verbose=False
+model_name = "samzito12/lora_model"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,
+    device_map="cpu",
+    low_cpu_mem_usage=True
 )
 
-# System prompt to fix identity issue
-SYSTEM_PROMPT = """You are a helpful AI assistant based on Meta's Llama-3.2-3B model, fine-tuned on the FineTome dataset. You are NOT ChatGPT and you are NOT made by OpenAI. You were created as part of a university machine learning project."""
+SYSTEM_PROMPT = "You are a helpful AI assistant based on Meta's Llama-3.2-3B model, fine-tuned on the FineTome dataset. You are NOT ChatGPT and NOT made by OpenAI."
 
 def chat(message, history):
-    """Generate response from YOUR fine-tuned model"""
+    # Build conversation
+    conversation = f"System: {SYSTEM_PROMPT}\n\n"
     
-    # Build conversation with system prompt
-    conversation = f"<|start_header_id|>system<|end_header_id|>\n\n{SYSTEM_PROMPT}<|eot_id|>"
-    
-    # Add chat history
     for user_msg, assistant_msg in history:
-        conversation += f"<|start_header_id|>user<|end_header_id|>\n\n{user_msg}<|eot_id|>"
-        conversation += f"<|start_header_id|>assistant<|end_header_id|>\n\n{assistant_msg}<|eot_id|>"
+        conversation += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
+    
+    conversation += f"User: {message}\nAssistant:"
     
-    # Add current message
-    conversation += f"<|start_header_id|>user<|end_header_id|>\n\n{message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+    # Tokenize
+    inputs = tokenizer(conversation, return_tensors="pt", truncate=True, max_length=2048)
     
-    # Generate response
-    response = llm(
-        conversation,
-        max_tokens=512,
-        temperature=0.7,
-        top_p=0.9,
-        stop=["<|eot_id|>", "<|start_header_id|>"],
-        echo=False
-    )
+    # Generate
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
     
-    return response['choices'][0]['text'].strip()
+    # Decode
+    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract assistant's response
+    if "Assistant:" in full_response:
+        response = full_response.split("Assistant:")[-1].strip()
+    else:
+        response = full_response[len(conversation):].strip()
+    
+    return response
 
-# Create Gradio interface
 demo = gr.ChatInterface(
     chat,
     title="🦙 My Fine-Tuned Llama-3.2-3B Chatbot",
     description="""
-    **Model**: Llama-3.2-3B fine-tuned on FineTome-100k dataset
-    
-    This chatbot uses a custom fine-tuned model, NOT ChatGPT.
+    **Model:** Llama-3.2-3B fine-tuned on FineTome-100k dataset
     
-    Created for ID2223 Lab 2 at KTH.
+    This is NOT ChatGPT - it's a custom fine-tuned model for ID2223 Lab 2.
     """,
     examples=[
         "What model are you?",
         "Explain machine learning in simple terms",
-        "Write a Python function to reverse a string",
-        "What is the weather like in Stockholm?"
+        "Write a Python function to reverse a string"
     ],
-    theme="soft",
-    retry_btn="🔄 Retry",
-    undo_btn="↩️ Undo",
-    clear_btn="🗑️ Clear"
+    theme="soft"
 )
 
 if __name__ == "__main__":