mejoras

app.py

@@ -1,50 +1,57 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-from huggingface_hub import InferenceClient
 from typing import List, Tuple, Dict
+import torch
 
-client = InferenceClient("AuriLab/gpt-bi-instruct-cesar")
+# Load model and tokenizer
+model_name = "AuriLab/gpt-bi-instruct-cesar"
+tokenizer_name = "AuriLab/gpt-bi"
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def format_messages(history: List[Tuple[str, str]], system_message: str, user_message: str) -> List[Dict[str, str]]:
-    messages = [{"role": "system", "content": system_message}]
-    messages.extend([
-        {"role": "user" if i % 2 == 0 else "assistant", "content": msg}
-        for turn in history
-        for i, msg in enumerate(turn)
-        if msg
-    ])
-    messages.append({"role": "user", "content": user_message})
-    return messages
+def format_messages(history: List[Tuple[str, str]], system_message: str, user_message: str) -> str:
+    # Format conversation history into a single string
+    formatted_prompt = system_message + "\n\n"
+    for user, assistant in history:
+        if user:
+            formatted_prompt += f"User: {user}\n"
+        if assistant:
+            formatted_prompt += f"Assistant: {assistant}\n"
+    formatted_prompt += f"User: {user_message}\nAssistant:"
+    return formatted_prompt
 
-def respond(message: str, history: List[Tuple[str, str]], system_message: str, max_tokens: int, temperature: float, top_p: float) -> str:
-    messages = format_messages(history, system_message, message)
-    response = ""
+def respond(message: str, history: List[Tuple[str, str]]) -> str:
+    system_message = """You are a helpful assistant. Follow these rules:
+                    1. Provide diverse and varied responses
+                    2. Avoid repeating the same words or phrases
+                    3. Use synonyms and alternative expressions
+                    4. Be concise and direct"""
     
-    for msg in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=0.7,        # Aumentado para más variedad
-        top_p=0.85,            # Ajustado para mejor balance
-    ):
-        token = msg.choices[0].delta.content
-        response += token
-        yield response
+    prompt = format_messages(history, system_message, message)
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    
+    # Generate response
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs["input_ids"],
+            max_new_tokens=200,
+            temperature=0.7,
+            top_p=0.85,
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Extract only the assistant's response
+    response = response.split("Assistant:")[-1].strip()
+    
+    return response
 
+# Create the Gradio interface with custom title
 demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(
-            value="""You are a helpful assistant. Follow these rules:
-1. Provide diverse and varied responses
-2. Avoid repeating the same words or phrases
-3. Use synonyms and alternative expressions
-4. Be concise and direct""",
-            label="System message"
-        ),
-        gr.Slider(minimum=1, maximum=256, value=200, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.7, maximum=1.2, value=1.0, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.85, step=0.05, label="Top-p (nucleus sampling)"),
-    ],
+    fn=respond,
+    title="Demo GPT-BI instruct",
 )
 
 if __name__ == "__main__":