Fix app.py: Replace InferenceClient with direct model loading using transformers

app.py

@@ -1,6 +1,18 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# Load model and tokenizer at startup
+model_name = "anktechsol/anki-2.5"
+print(f"Loading model {model_name}...")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float32,
+    device_map="auto",
+    low_cpu_mem_usage=True
+)
+print("Model loaded successfully!")
 
 def respond(
     message,
@@ -9,36 +21,75 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-    hf_token: gr.OAuthToken,
 ):
     """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+    Generate responses using the local Anki 2.5 model.
     """
-    client = InferenceClient(token=hf_token.token, model="anktechsol/anki-2.5")
-
-    messages = [{"role": "system", "content": system_message}]
-
-    messages.extend(history)
-
-    messages.append({"role": "user", "content": message})
-
+    # Build conversation history
+    conversation = []
+    
+    # Add system message if provided
+    if system_message:
+        conversation.append({"role": "system", "content": system_message})
+    
+    # Add chat history
+    for msg in history:
+        conversation.append(msg)
+    
+    # Add current message
+    conversation.append({"role": "user", "content": message})
+    
+    # Format prompt for the model
+    # Try to apply chat template if available, otherwise use simple format
+    try:
+        formatted_prompt = tokenizer.apply_chat_template(
+            conversation,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+    except:
+        # Fallback to simple format if chat template not available
+        formatted_prompt = ""
+        for msg in conversation:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            if role == "system":
+                formatted_prompt += f"System: {content}\n"
+            elif role == "user":
+                formatted_prompt += f"User: {content}\n"
+            elif role == "assistant":
+                formatted_prompt += f"Assistant: {content}\n"
+        formatted_prompt += "Assistant: "
+    
+    # Tokenize input
+    inputs = tokenizer(formatted_prompt, return_tensors="pt", truncation=True, max_length=2048)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    
+    # Generate response with streaming
     response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
-        yield response
-
+    with torch.no_grad():
+        for i in range(max_tokens):
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=1,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+            )
+            
+            # Decode only the new token
+            new_token = tokenizer.decode(outputs[0][-1], skip_special_tokens=True)
+            
+            # Check for end of generation
+            if outputs[0][-1].item() == tokenizer.eos_token_id:
+                break
+            
+            response += new_token
+            yield response
+            
+            # Update inputs for next iteration
+            inputs = {"input_ids": outputs, "attention_mask": torch.ones_like(outputs)}
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -47,9 +98,9 @@ chatbot = gr.ChatInterface(
     respond,
     type="messages",
     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.Textbox(value="You are a friendly Chatbot proficient in Indian languages.", label="System message"),
+        gr.Slider(minimum=1, maximum=512, value=256, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
             maximum=1.0,
@@ -60,11 +111,5 @@ chatbot = gr.ChatInterface(
     ],
 )
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
-
-
 if __name__ == "__main__":
-    demo.launch()
+    chatbot.launch()