Update app.py

app.py

@@ -1,50 +1,95 @@
-import os
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
-
-# 1. Autenticación
-HF_TOKEN = os.environ["HF_TOKEN"]
-os.environ["HUGGINGFACEHUB_API_TOKEN"] = HF_TOKEN
-
-# 2. Carga del modelo y tokenizer en GPU
-MODEL_ID = "arcee-ai/AFM-4.5B"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
-model     = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    use_auth_token=HF_TOKEN,
-    torch_dtype=torch.bfloat16,
-    device_map="auto"
-)
-chat = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device_map="auto"
+import requests
+import json
+import os
+
+"""
+Using Together AI API for chat completions
+"""
+TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")
+if not TOGETHER_API_KEY:
+    raise ValueError("TOGETHER_API_KEY environment variable is not set")
+
+TOGETHER_API_URL = "https://api.together.xyz/v1/chat/completions"
+
+
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    messages = [{"role": "system", "content": system_message}]
+
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+
+    messages.append({"role": "user", "content": message})
+
+    headers = {
+        "Authorization": f"Bearer {TOGETHER_API_KEY}",
+        "Content-Type": "application/json"
+    }
+    
+    data = {
+        "model": "arcee-ai/AFM-4.5B",
+        "messages": messages,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "stream": True
+    }
+
+    response = ""
+    
+    try:
+        with requests.post(TOGETHER_API_URL, headers=headers, json=data, stream=True) as r:
+            r.raise_for_status()
+            for line in r.iter_lines():
+                if line:
+                    line = line.decode('utf-8')
+                    if line.startswith('data: '):
+                        line = line[6:]  # Remove 'data: ' prefix
+                        if line.strip() == '[DONE]':
+                            break
+                        try:
+                            chunk = json.loads(line)
+                            if 'choices' in chunk and len(chunk['choices']) > 0:
+                                delta = chunk['choices'][0].get('delta', {})
+                                if 'content' in delta:
+                                    token = delta['content']
+                                    response += token
+                                    yield response
+                        except json.JSONDecodeError:
+                            continue
+    except requests.exceptions.RequestException as e:
+        yield f"Error: {str(e)}"
+
+
+"""
+For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+"""
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
+    ],
 )
 
-# 3. Función de inferencia
-def genera_respuesta(prompt: str, max_tokens: int = 256, temp: float = 0.5, top_p: float = 0.95):
-    out = chat(
-        prompt,
-        max_new_tokens=max_tokens,
-        temperature=temp,
-        top_p=top_p,
-        do_sample=True
-    )
-    return out[0]["generated_text"].strip()
-
-# 4. Interfaz Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("### AFM-4.5B en tu Space")
-    with gr.Row():
-        inp = gr.Textbox(label="Pregunta", lines=2)
-        out = gr.Textbox(label="Respuesta")
-    with gr.Row():
-        max_toks = gr.Slider(50, 512, value=256, label="Max new tokens")
-        temp    = gr.Slider(0.1, 1.0, value=0.5, label="Temperatura")
-        top_p   = gr.Slider(0.1, 1.0, value=0.95, label="Top-p")
-    btn = gr.Button("Enviar")
-    btn.click(fn=genera_respuesta, inputs=[inp, max_toks, temp, top_p], outputs=out)
-
-demo.launch()
+
+if __name__ == "__main__":
+    demo.launch()