Update app.py

app.py

@@ -1,74 +1,105 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+)
 import os
+from threading import Thread
+import spaces
+import time
+import subprocess
 
-# Model and tokenizer paths
-model_path = "rajj0/autotrain-phi3-midium-4k-godsent-orpo-6"
-hf_token = os.getenv("HF_TOKEN")  # Get the token from environment variables
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
 
-# Debugging: print the token to ensure it's being set
-print(f"HF_TOKEN: {hf_token}")
+token = os.environ["HF_TOKEN"]
 
-if hf_token is None:
-    raise ValueError("HF_TOKEN environment variable not set")
 
-try:
-    # Load the tokenizer and model with trust_remote_code=True
-    print("Loading tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(model_path, use_auth_token=hf_token, trust_remote_code=True)
-    print("Tokenizer loaded successfully.")
+model = AutoModelForCausalLM.from_pretrained(
+    "rajj0/autotrain-phi3-midium-4k-godsent-orpo-6",
+    token=token,
+    trust_remote_code=True,
+)
+tok = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct", token=token)
+terminators = [
+    tok.eos_token_id,
+]
 
-    print("Loading model...")
-    model = AutoModelForCausalLM.from_pretrained(
-        model_path,
-        device_map="auto",
-        torch_dtype='auto',
-        use_auth_token=hf_token,
-        trust_remote_code=True
-    ).eval()
-    print("Model loaded successfully.")
-except Exception as e:
-    print(f"Error loading model or tokenizer: {e}")
-    raise
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print(f"Using GPU: {torch.cuda.get_device_name(device)}")
+else:
+    device = torch.device("cpu")
+    print("Using CPU")
 
-# Function to generate a response from the model
-def generate_response(user_input):
-    try:
-        print(f"User input: {user_input}")
-        messages = [{"role": "user", "content": user_input}]
-        print(f"Messages: {messages}")
+model = model.to(device)
+# Dispatch Errors
 
-        input_ids = tokenizer.apply_chat_template(
-            conversation=messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors='pt'
-        )
-        print(f"Input IDs: {input_ids}")
 
-        output_ids = model.generate(input_ids.to('cuda'))
-        print(f"Output IDs: {output_ids}")
+@spaces.GPU(duration=60)
+def chat(message, history, temperature, do_sample, max_tokens):
+    chat = []
+    for item in history:
+        chat.append({"role": "user", "content": item[0]})
+        if item[1] is not None:
+            chat.append({"role": "assistant", "content": item[1]})
+    chat.append({"role": "user", "content": message})
+    messages = tok.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+    model_inputs = tok([messages], return_tensors="pt").to(device)
+    streamer = TextIteratorStreamer(
+        tok, timeout=20.0, skip_prompt=True, skip_special_tokens=True
+    )
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        eos_token_id=terminators,
+    )
 
-        response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
-        print(f"Response: {response}")
+    if temperature == 0:
+        generate_kwargs["do_sample"] = False
 
-        return response
-    except Exception as e:
-        print(f"Error generating response: {e}")
-        return "An error occurred while generating the response."
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
 
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=generate_response,
-    inputs="text",
-    outputs="text",
-    title="PHI Model Chatbot",
-    description="A chatbot powered by the PHI model."
-)
+    partial_text = ""
+    for new_text in streamer:
+        partial_text += new_text
+        yield partial_text
 
-# Launch the Gradio interface
-if __name__ == "__main__":
-    print("Launching Gradio interface...")
-    iface.launch()
-    print("Gradio interface launched.")
+    yield partial_text
+
+
+demo = gr.ChatInterface(
+    fn=chat,
+    examples=[["Write me a poem about Machine Learning."]],
+    # multimodal=False,
+    additional_inputs_accordion=gr.Accordion(
+        label="⚙️ Parameters", open=False, render=False
+    ),
+    additional_inputs=[
+        gr.Slider(
+            minimum=0, maximum=1, step=0.1, value=0.9, label="Temperature", render=False
+        ),
+        gr.Checkbox(label="Sampling", value=True),
+        gr.Slider(
+            minimum=128,
+            maximum=4096,
+            step=1,
+            value=512,
+            label="Max new tokens",
+            render=False,
+        ),
+    ],
+    stop_btn="Stop Generation",
+    title="Chat With LLMs",
+    description="Now Running [microsoft/Phi-3-mini-128k-instruct](https://huggingface.co/microsoft/Phi-3-mini-128k-instruct)",
+)
+demo.launch()