Update app.py

app.py

@@ -1,35 +1,97 @@
-import torch 
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline 
-
-torch.random.manual_seed(0) 
-model = AutoModelForCausalLM.from_pretrained( 
-    "microsoft/Phi-3-mini-4k-instruct",  
-    device_map="cuda",  
-    torch_dtype="auto",  
-    trust_remote_code=True,  
-) 
-
-tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct") 
-
-messages = [ 
-    {"role": "system", "content": "You are a helpful AI assistant."}, 
-    {"role": "user", "content": "Can you provide ways to eat combinations of bananas and dragonfruits?"}, 
-    {"role": "assistant", "content": "Sure! Here are some ways to eat bananas and dragonfruits together: 1. Banana and dragonfruit smoothie: Blend bananas and dragonfruits together with some milk and honey. 2. Banana and dragonfruit salad: Mix sliced bananas and dragonfruits together with some lemon juice and honey."}, 
-    {"role": "user", "content": "What about solving an 2x + 3 = 7 equation?"}, 
-] 
-
-pipe = pipeline( 
-    "text-generation", 
-    model=model, 
-    tokenizer=tokenizer, 
-) 
-
-generation_args = { 
-    "max_new_tokens": 500, 
-    "return_full_text": False, 
-    "temperature": 0.0, 
-    "do_sample": False, 
-} 
-
-output = pipe(messages, **generation_args) 
-print(output[0]['generated_text'])
+import gradio as gr
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+)
+import os
+from threading import Thread
+import spaces
+import time
+
+token = os.environ["HF_TOKEN"]
+
+
+model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/Phi-3-mini-4k-instruct", token=token,trust_remote_code=True
+)
+tok = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct", token=token)
+terminators = [
+    tok.eos_token_id,
+]
+
+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")
+
+model = model.to(device)
+# Dispatch Errors
+
+
+@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,
+    )
+    
+    if temperature == 0:
+        generate_kwargs['do_sample'] = False
+    
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    partial_text = ""
+    for new_text in streamer:
+        partial_text += new_text
+        yield partial_text
+
+
+    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-4k-instruct](https://huggingface.com/microsoft/Phi-3-mini-4k-instruct)"
+)
+demo.launch()