Update app.py

app.py

@@ -1,22 +1,40 @@
 import os
 import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel, PeftConfig
 
-# Constants for generation parameters
 MAX_NEW_TOKENS = 100
 TEMPERATURE = 0.5
 TOP_P = 0.95
 TOP_K = 50
 REPETITION_PENALTY = 1.05
 
-
-# Global variables to store model and tokenizer
-model = None
-tokenizer = None
+HF_TOKEN = os.getenv('HF_TOKEN')
 
 def load_model():
-    global model, tokenizer
+    base_model_id = "meta-llama/Llama-2-7b-hf"
+    peft_model_id = "somosnlp-hackathon-2025/Llama-2-7b-hf-lora-refranes"
+
+    config = PeftConfig.from_pretrained(peft_model_id)
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_model_id,
+        torch_dtype="auto",
+        device_map="auto",
+        token=HF_TOKEN
+    )
+
+    model = PeftModel.from_pretrained(base_model, peft_model_id)
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
     return model, tokenizer
 
+model = None
+tokenizer = None
 
 def generate_response(input_text, max_tokens, temperature, top_p, repetition_penalty):
     global model, tokenizer
@@ -24,23 +42,41 @@ def generate_response(input_text, max_tokens, temperature, top_p, repetition_pen
     if model is None or tokenizer is None:
         model, tokenizer = load_model()
     
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
     
-    return "Adios"
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            do_sample=True,
+            top_p=top_p,
+            top_k=TOP_K,
+            repetition_penalty=repetition_penalty
+        )
 
-def chat_interface(message, history, system_message, max_tokens, temperature, top_p, repetition_penalty):
+    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    if "->" in full_response:
+        response_parts = full_response.split("->", 1)
+        if len(response_parts) > 1:
+            return response_parts[1].strip()
+    
+    return full_response.strip()
 
-  prompt = f"{message}"
-  if system_message:
-      prompt = f"{system_message}\n{message}"
-      
-  response = generate_response(
-      prompt, 
-      max_tokens, 
-      temperature, 
-      top_p, 
-      repetition_penalty
-  )
-  return response
+def chat_interface(message, history, system_message, max_tokens, temperature, top_p, repetition_penalty):
+    prompt = f"{message}"
+    if system_message:
+        prompt = f"{system_message}\n{message}"
+        
+    response = generate_response(
+        prompt, 
+        max_tokens, 
+        temperature, 
+        top_p, 
+        repetition_penalty
+    )
+    return response
 
 demo = gr.ChatInterface(
     chat_interface,