Update app.py

app.py

@@ -1,31 +1,64 @@
+import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
-# Load the finetuned model and tokenizer from Hugging Face Model Hub
-model_path = "sagar007/phi3.5_finetune"
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device_map="auto")
-
-# Create a text-generation pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-def generate_text(prompt, max_length=100, temperature=0.7):
-    """Generate text based on the input prompt."""
-    generated = generator(prompt, max_length=max_length, temperature=temperature, num_return_sequences=1)
-    return generated[0]['generated_text']
-
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.Textbox(lines=5, label="Enter your prompt"),
-        gr.Slider(minimum=50, maximum=500, value=100, step=10, label="Max Length"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
+import spaces
+
+# Check if CUDA is available and set the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Load model and tokenizer
+MODEL_PATH = "sagar007/phi3.5_finetune"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/Phi-3.5-mini-instruct",
+    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+peft_config = PeftConfig.from_pretrained(MODEL_PATH)
+model = PeftModel.from_pretrained(base_model, MODEL_PATH)
+model.to(device)
+model.eval()
+
+@spaces.GPU(duration=60)
+def generate_response(instruction, max_length=512):
+    prompt = f"Instruction: {instruction}\nResponse:"
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response.split("Response:")[1].strip()
+
+def chatbot(message, history):
+    response = generate_response(message)
+    return response
+
+demo = gr.ChatInterface(
+    chatbot,
+    title="Fine-tuned Phi-3.5 Chatbot",
+    description="This is a chatbot using a fine-tuned version of the Phi-2 model.",
+    theme="default",
+    examples=[
+        "Explain the concept of machine learning.",
+        "Write a short story about a robot learning to paint.",
+        "What are some effective ways to reduce stress?",
     ],
-    outputs=gr.Textbox(lines=10, label="Generated Text"),
-    title="Finetuned Phi-3.5 Text Generation",
-    description="Enter a prompt and generate text using the finetuned Phi-3.5 model.",
+    cache_examples=True,
 )
 
-# Launch the app
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()