Update main.py

main.py

@@ -2,34 +2,35 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-import os
+from hugchat import hugchat
+from hugchat.login import Login
 import asyncio
+import os
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
 
-# Set cache directories
-cache_dir = "/tmp/hf_home"
-os.environ["HF_HOME"] = cache_dir
-os.environ["TRANSFORMERS_CACHE"] = cache_dir
-os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
+# Read credentials from environment variables
+EMAIL = os.getenv("EMAIL")
+PASSWD = os.getenv("PASSWD")
 
-os.makedirs(cache_dir, exist_ok=True)
-os.chmod(cache_dir, 0o777)
+# Cookie storage
+cookie_path_dir = "./cookies/"
+os.makedirs(cookie_path_dir, exist_ok=True)
 
-# Load model and tokenizer
-model_name = "EleutherAI/gpt-neo-1.3B"
-tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
-model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)
+# HugChat login
+sign = Login(EMAIL, PASSWD)
+cookies = sign.login(cookie_dir_path=cookie_path_dir, save_cookies=True)
 
-# Set pad token if not defined
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+# Create chatbot instance
+chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
 
-# Set device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+# Optional: Use assistant ID
+ASSISTANT_ID = "66017fca58d60bd7d5c5c26c"  # Replace if needed
+chatbot.new_conversation(assistant=ASSISTANT_ID, switch_to=True)
 
-# Initialize FastAPI
+# FastAPI setup
 app = FastAPI()
 
 # Enable CORS
@@ -41,57 +42,22 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Request model
 class Question(BaseModel):
     question: str
 
-SYSTEM_PROMPT = "You are a helpful, professional, and highly persuasive sales assistant..."
-
-chat_history_ids = None
-
-async def generate_response_chunks(prompt: str):
-    global chat_history_ids
-
-    # Combine system prompt and user input
-    input_text = SYSTEM_PROMPT + "\nUser: " + prompt + "\nBot:"
-    new_input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)
-    
-    # Create attention mask (handle case where pad_token_id might be None)
-    attention_mask = torch.ones_like(new_input_ids)
-    
-    if chat_history_ids is not None:
-        input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
-        attention_mask = torch.cat([
-            torch.ones_like(chat_history_ids),
-            attention_mask
-        ], dim=-1)
-    else:
-        input_ids = new_input_ids
-
-    # Generate response
-    output_ids = model.generate(
-        input_ids,
-        attention_mask=attention_mask,
-        max_new_tokens=200,
-        do_sample=True,
-        top_p=0.9,
-        temperature=0.7,
-        pad_token_id=tokenizer.eos_token_id
-    )
-
-    # Update chat history
-    chat_history_ids = output_ids
-
-    # Decode only the new tokens
-    response = tokenizer.decode(output_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-
-    # Stream the response
-    for word in response.split():
-        yield word + " "
-        await asyncio.sleep(0.03)
+# Token stream function
+async def generate_response_stream(prompt: str):
+    for chunk in chatbot.chat(prompt, stream=True):
+        token = chunk.get("token", "")
+        if token:
+            yield token
+            await asyncio.sleep(0.02)
 
+# Endpoint
 @app.post("/ask")
 async def ask(question: Question):
     return StreamingResponse(
-        generate_response_chunks(question.question),
+        generate_response_stream(question.question),
         media_type="text/plain"
-    )
+    )