Update app.py

app.py

@@ -1,8 +1,18 @@
 import gradio as gr
 import torch
-import torch.nn.functional as F
+import time
+import json
+import uuid
+import os
+import pytz
+from datetime import datetime
 from transformers import AutoTokenizer
 from unsloth import FastLanguageModel
+from pathlib import Path
+from huggingface_hub import CommitScheduler
+
+# Load HF token from the environment
+token = os.environ["HF_TOKEN"]
 
 # Model Setup
 max_seq_length = 2048
@@ -28,14 +38,45 @@ model = FastLanguageModel.get_peft_model(
 )
 FastLanguageModel.for_inference(model)
 
-def generate_response(task, input_text, temperature, top_p, max_tokens):
-    prompt = f"### INPUT : {input_text} RESPONSE : "
-    message = [{"role": "user", "content": prompt}]
-    inputs = tokenizer.apply_chat_template(message, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")
+# Task-Specific Prompt Mapping
+option_mapping = {
+    "translation": "### TRANSLATION ###",
+    "mcq": "### MCQ ###",
+    "nli": "### NLI ###",
+    "summarization": "### SUMMARIZATION ###",
+    "long response": "### LONG RESPONSE ###",
+    "direct response": "### DIRECT RESPONSE ###",
+    "paraphrase": "### PARAPHRASE ###",
+    "code": "### CODE ###",
+}
+
+# Set up logging folder and CommitScheduler
+log_folder = Path("logs")
+log_folder.mkdir(parents=True, exist_ok=True)
+log_file = log_folder / f"chat_log_{uuid.uuid4()}.json"
+
+scheduler = CommitScheduler(
+    repo_id="DrishtiSharma/phi-4-unsloth-logs",  
+    repo_type="dataset",
+    folder_path=log_folder,
+    path_in_repo="data",
+    every=10,
+    token=token
+)
+
+# Fixed timezone
+timezone = pytz.timezone("UTC")
+
+def generate_response(message, temperature, max_new_tokens, top_p, task):
+    append_text = option_mapping.get(task, "")
+    prompt = f"### INPUT : {message} {append_text} RESPONSE : "
+    print(f"Prompt: {prompt}")
+    start_time = time.time()
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
     
     outputs = model.generate(
         input_ids=inputs,
-        max_new_tokens=max_tokens,
+        max_new_tokens=max_new_tokens,
         use_cache=True,
         temperature=temperature,
         top_p=top_p,
@@ -43,30 +84,50 @@ def generate_response(task, input_text, temperature, top_p, max_tokens):
     )
     
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    processed_response = response.split("### RESPONSE :assistant")[-1].strip()
+    processed_response = response.split("### RESPONSE :")[-1].strip()
+    end_time = time.time()
+    
+    response_time = round(end_time - start_time, 2)
+    timestamp = datetime.now(timezone).strftime("%Y-%m-%d %H:%M:%S %Z")
+    log_data = {
+        "timestamp": timestamp,
+        "input": message,
+        "output": processed_response,
+        "response_time": response_time,
+        "temperature": temperature,
+        "max_tokens": max_new_tokens,
+        "top_p": top_p
+    }
+    
+    with scheduler.lock:
+        with log_file.open("a") as f:
+            f.write(json.dumps(log_data) + "\n")
+    
     return processed_response
 
-# Gradio Interface
-def gradio_ui():
-    with gr.Blocks() as demo:
-        gr.Markdown("## Test Space: Chat with Phi-4-Hindi")
-        with gr.Row():
-            task = gr.Dropdown([
-                "Long Response", "Short Response", "NLI", "Translation", "MCQ", "Cross-Lingual"
-            ], label="Select Task")
-            input_text = gr.Textbox(label="Input Text")
-        
-        with gr.Row():
-            temperature = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature")
-            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.1, label="Top P")
-            max_tokens = gr.Slider(50, 800, value=200, step=50, label="Max Tokens")
-        
-        output_text = gr.Textbox(label="Generated Response")
-        btn = gr.Button("Generate")
-        btn.click(generate_response, inputs=[task, input_text, temperature, top_p, max_tokens], outputs=output_text)
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## Chat with Phi-4-Hindi")
+    
+    task_dropdown = gr.Dropdown(
+        choices=list(option_mapping.keys()),
+        value="long response",
+        label="Select Task"
+    )
+    message_input = gr.Textbox(label="Enter your message")
     
-    return demo
+    with gr.Row():
+        temperature_slider = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature")
+        top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.1, label="Top P")
+        max_tokens_slider = gr.Slider(50, 800, value=200, step=50, label="Max Tokens")
+    
+    output_box = gr.Textbox(label="Generated Response")
+    generate_btn = gr.Button("Generate")
+    
+    generate_btn.click(
+        generate_response,
+        inputs=[message_input, temperature_slider, max_tokens_slider, top_p_slider, task_dropdown],
+        outputs=output_box
+    )
 
-# Launch Gradio App
-demo = gradio_ui()
 demo.launch()